Question about jax-metal development and missing operations

[a-gn]

Hi! I've tested jax-metal on a mac and encountered a few of these missing operations, and others, while implementing common papers.

I'd like to know if there is an ETA for ~full support of Apple Silicon platforms in JAX. It would be a great help if the neural/GPU cores on my mac could replace cloud compute for some of these workflows, but right now I have to convert to numpy and back for some operations.

I'm also interested in contributing, but not sure if it's within my competence yet. Is there something I can read (part of the codebase, guide, etc.) to know where to start adding support for new operations on Metal?

Thanks!

[jakevdp]

cc/ @shuhand0 @kulinseth

[thegodone]

I see that the CI test are failing since 4 months...

Any idea why Apple plug in is not updated too ?

[shuhand0]

The failures on the CI runner are due to incompatibility between jaxlib and jax-metal. We have been working with JAX team to address it and are going to make a patch release soon.

[jakiw]

Thanks for the great library, I am using it every day.
However, some relatively important features are still missing (especially debug.print and cholesky facotrization). I have seen there have been no updates to the library since last summer, whereas it has been updated quite frequently before.

Has the library been abandoned for now // been given a very low priority? Or are updates planned?

[jakevdp]

jax-metal is closed-source, and the JAX team has no access to the codebase, so unfortunately we don't have any influence on the timeline or roadmap for jax-metal support. The Apple engineers who maintain jax-metal are tagged here (@shuhand0 & @kulinseth), so hopefully they'll be able to comment.