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.