Open yu-fz opened 1 year ago
Just to add this and a plethora of other fundamental jax operations simply don't work on Jax-metal with no update insight. At this point, jax-metal is quite useless as far as I have been able to test it. As of version Jax-metal 0.0.4, I would say simply avoid it and install Jax on CPU.
Just FYI, it seems like Apple released a new version of metal-jax a few weeks back: version 0.0.5. This basic feature is still broken.
Same issue here. Any idea when a fix will be released?
I will ping this and add that I hope this gets fixed in the next release.
I think the JAX + Apple Silicon combo offers some unique advantages for prototyping models locally (i.e. smaller compute but larger RAM compared to a comparable NVIDIA workstation). I think JAX users will pick up on this pretty quickly once features like this are fixed.
We are looking into adding the conversion for the op.
I also hope this gets fixed soon! Do you happen to have any updates @shuhand0? Thanks
As of jax-metal 0.1.1 this issue still persists, at least in my configuration.
MacOS 15.0.1 Apple M2 Pro
As of jax-metal 0.1.1 this issue still persists, at least in my configuration.
MacOS 15.0.1 Apple M2 Pro
😭 😭 😭 bruh
Description
Matrix inversion appears to be broken on jax-metal. Apologies in advance if this is not the right place to report the issue.
Repro:
What jax/jaxlib version are you using?
jax 0.4.11, jaxlib 0.4.10
Which accelerator(s) are you using?
GPU/Metal
Additional system info
Mac OS 13.5.1, M1 Max
NVIDIA GPU info
No response