Jax-Metal encounters xll_extension runtime error

[acranej]

Hello,

I am running jax-metal version 0.0.3 with jaxlib 0.4.10 as per Jax-Metal. This is on an Apple M1 Pro in a clean virtual environment following the instructions in the link above.
I am getting the following error: "jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: :0: error: 'func.func' op One or more function input/output data types are not supported.". I have tried to specifically dictate float32 in pytensor using the following: pytensor.config.floatX="float32". However, the error stays. Is there something I am missing? It seemed like Jax-Metal was supported for M1 GPU usage.

[jakevdp]

cc/ @shuhand0 @kulinseth

[adam-hartshorne]

jax-metal 0.0.3 is very old. The latest version is 0.1.0

[acranej]

I'm now running jax-metal v0.1.0, jax v0.4.26, jaxlib v0.4.26, and numpy v1.26.4. I am recieving this error:

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Unable to serialize MPS module

[jakevdp]

This may be due to an existing issue, but it's difficult to say because you didn't share a reproduction. See Issues with the jax-metal tag for a list of known issues.

[jakevdp]

Actually, this looks like your issue: #20401

[acranej]

I did see this as well, I got rid of the error but now have a segmentation fault. I can try to make a reproducible example but may be out of the scope of this thread. Thanks for the help!