Is anyone working on jax-metal?

Hi,

I think many of us would love to be able to use our GPUs for Jax on the new Apple Silicon devices, but currently, the Jax-metal plugin is, for all effects and purposes, broken. Is it still under active development? Is there a planned release for a new version?

thanks!

Another interested user here. The bad news is I think there is still lots of things need to be done before Jax-metal can be well use

I'm also interested in this.

Apple Silicon is an exciting hardware platform for ML and we need guidance from Apple how it envisions the future ecosystem. Is there going to be more openness and support of the major ML frameworks (incl. JAX), or is jax-metal an abandoned tech demo and we're all meant to use MLX more or less exclusively?
The latter scenario would be incredibly unfortunate.

I'm also interested. JAX can also be used for other purposes that MLX cannot. Currently for robotics simulation MuJoCo MJX uses JAX. Unfortunately jax-metal is still broken for that to be run MacOS.

Is anyone working on jax-metal?
 
 
Q