Any chance of a bit of support for jax-metal, or incorporating apple silicon support into Jax?