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.