Unable to use Jax with metal on Apple M2

Question

Created Jul ’23

Replies 2

Boosts 0

Participants 2

Hi, I'm following https://developer.apple.com/metal/jax/ to install jax on my Mac. The installation is successful. However, running the give example gives

$ python -c 'import jax; jax.numpy.arange(10)'
2023-07-27 20:26:08.492162: W pjrt_plugin/src/mps_client.cc:535] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported!
Metal device set to: Apple M2 Pro

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

loc("-":2:3): error: custom op 'func.func' is unknown
fish: Job 1, 'python3 $argv' terminated by signal SIGSEGV (Address boundary error)

Boost

Answer 1

fmarquardt OP

Aug ’23

The same just happened to me on M1.

I had followed the instructions to build jaxlib and jax, and this is what I get:

Python 3.9.6 | packaged by conda-forge | (default, Jul 11 2021, 03:35:11) 
[Clang 11.1.0 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.

import jax.numpy as jnp

jnp.linspace(-1,1,100)

2023-08-12 15:34:46.663492: W pjrt_plugin/src/mps_client.cc:535] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported!
Metal device set to: Apple M1

systemMemory: 16.00 GB

maxCacheSize: 5.33 GB

loc("-":2:3): error: custom op 'func.func' is unknown
Segmentation fault: 11

0

Answer 2

fmarquardt OP

Aug ’23

I am happy to report that the issue went away for me after I upgraded from Monterey to the latest Ventura (13.5) OS --- without changing anything else in my installation.

Probably I should have read the instructions on the metal/jax Apple page more carefully, where they say one should have at least 13.4 .

0