Agreed, I would love to use jax-metal, but it is half a year out of date and pip cannot find any version of flax that works with this old version.