But it crashes Python when I try to run a compiled function. And that's only after discovering I need an older specific version of jax-metal, because newer versions apparently don't work with M1 anymore (only M2/M3) -- they don't even report the GPU as existing. And even if you get it running, it's missing support for complex numbers.
I'm not clear whether it's Google or Apple who is building/maintaining support for Apple M chips though.
JAX works perfectly in CPU mode though on my MBA, so at least I can use it for development and debugging.
Pretty sure it's Apple building it, and they're using JAX in-house so I imagine it will get better over time. Though they do love to drop support for old things so maybe M1 will never work again...
I think they’re likely using MLX in house now, no? (Probably not everyone, ofc - but seems likely that many will just use the native array framework designed explicitly for MX chips)
Sadly, I cannot get JAX to work with the built-in GPU on my M1 MacBook Air. In theory it's supposed to work:
https://developer.apple.com/metal/jax/
But it crashes Python when I try to run a compiled function. And that's only after discovering I need an older specific version of jax-metal, because newer versions apparently don't work with M1 anymore (only M2/M3) -- they don't even report the GPU as existing. And even if you get it running, it's missing support for complex numbers.
I'm not clear whether it's Google or Apple who is building/maintaining support for Apple M chips though.
JAX works perfectly in CPU mode though on my MBA, so at least I can use it for development and debugging.