jax-metal: Installation according to apples guide does not work.

[PhilippFoerster]

Description

Description: I am trying to install Jax with support for Apple Silicon.
I follow Apple's guide: https://developer.apple.com/metal/jax/

Install:

(base) philipp@Mac ~ % python3 -m venv ~/jax-metal
source ~/jax-metal/bin/activate
python -m pip install -U pip
python -m pip install numpy wheel
python -m pip install jax-metal
Requirement already satisfied: pip in ./jax-metal/lib/python3.13/site-packages (24.3.1)
Collecting pip
  Using cached pip-25.0.1-py3-none-any.whl.metadata (3.7 kB)
Using cached pip-25.0.1-py3-none-any.whl (1.8 MB)
Installing collected packages: pip
  Attempting uninstall: pip
    Found existing installation: pip 24.3.1
    Uninstalling pip-24.3.1:
      Successfully uninstalled pip-24.3.1
Successfully installed pip-25.0.1
Collecting numpy
  Using cached numpy-2.2.3-cp313-cp313-macosx_14_0_arm64.whl.metadata (62 kB)
Collecting wheel
  Downloading wheel-0.45.1-py3-none-any.whl.metadata (2.3 kB)
Using cached numpy-2.2.3-cp313-cp313-macosx_14_0_arm64.whl (5.1 MB)
Downloading wheel-0.45.1-py3-none-any.whl (72 kB)
Installing collected packages: wheel, numpy
Successfully installed numpy-2.2.3 wheel-0.45.1
Collecting jax-metal
  Using cached jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl.metadata (2.7 kB)
Requirement already satisfied: wheel~=0.35 in ./jax-metal/lib/python3.13/site-packages (from jax-metal) (0.45.1)
Collecting six>=1.15.0 (from jax-metal)
  Using cached six-1.17.0-py2.py3-none-any.whl.metadata (1.7 kB)
Collecting jaxlib>=0.4.34 (from jax-metal)
  Using cached jaxlib-0.5.1-cp313-cp313-macosx_11_0_arm64.whl.metadata (978 bytes)
Collecting jax>=0.4.34 (from jax-metal)
  Using cached jax-0.5.2-py3-none-any.whl.metadata (22 kB)
Collecting ml_dtypes>=0.4.0 (from jax>=0.4.34->jax-metal)
  Using cached ml_dtypes-0.5.1-cp313-cp313-macosx_10_13_universal2.whl.metadata (21 kB)
Requirement already satisfied: numpy>=1.25 in ./jax-metal/lib/python3.13/site-packages (from jax>=0.4.34->jax-metal) (2.2.3)
Collecting opt_einsum (from jax>=0.4.34->jax-metal)
  Using cached opt_einsum-3.4.0-py3-none-any.whl.metadata (6.3 kB)
Collecting scipy>=1.11.1 (from jax>=0.4.34->jax-metal)
  Using cached scipy-1.15.2-cp313-cp313-macosx_14_0_arm64.whl.metadata (61 kB)
Using cached jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl (41.2 MB)
Using cached jax-0.5.2-py3-none-any.whl (2.4 MB)
Using cached jaxlib-0.5.1-cp313-cp313-macosx_11_0_arm64.whl (81.1 MB)
Using cached six-1.17.0-py2.py3-none-any.whl (11 kB)
Using cached ml_dtypes-0.5.1-cp313-cp313-macosx_10_13_universal2.whl (667 kB)
Using cached scipy-1.15.2-cp313-cp313-macosx_14_0_arm64.whl (22.4 MB)
Using cached opt_einsum-3.4.0-py3-none-any.whl (71 kB)
Installing collected packages: six, scipy, opt_einsum, ml_dtypes, jaxlib, jax, jax-metal
Successfully installed jax-0.5.2 jax-metal-0.1.1 jaxlib-0.5.1 ml_dtypes-0.5.1 opt_einsum-3.4.0 scipy-1.15.2 six-1.17.0

Test:

(jax-metal) (base) philipp@Mac ~ % python -c 'import jax; print(jax.numpy.arange(10))'
WARNING:2025-03-06 16:36:46,318:jax._src.xla_bridge:997: Platform 'METAL' is experimental and not all JAX functionality may be correctly supported!
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1741275406.318312 13032203 mps_client.cc:510] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported!
Metal device set to: Apple M3

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

I0000 00:00:1741275406.327869 13032203 service.cc:145] XLA service 0x600002c9ce00 initialized for platform METAL (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1741275406.327885 13032203 service.cc:153]   StreamExecutor device (0): Metal, <undefined>
I0000 00:00:1741275406.329246 13032203 mps_client.cc:406] Using Simple allocator.
I0000 00:00:1741275406.329252 13032203 mps_client.cc:384] XLA backend will use up to 11452776448 bytes on device 0 for SimpleAllocator.
Traceback (most recent call last):
  File "<string>", line 1, in <module>
    import jax; print(jax.numpy.arange(10))
                      ~~~~~~~~~~~~~~~~^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py", line 6317, in arange
    output = _arange(start, stop=stop, step=step, dtype=dtype)
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py", line 6357, in _arange
    return lax.iota(dtype, start)  # type: ignore[arg-type]
           ~~~~~~~~^^^^^^^^^^^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/lax/lax.py", line 2932, in iota
    return broadcasted_iota(dtype, (size,), 0)
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/lax/lax.py", line 2944, in broadcasted_iota
    return iota_p.bind(*dynamic_shape, dtype=dtype, shape=tuple(static_shape),
           ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                       dimension=dimension, sharding=out_sharding)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/core.py", line 502, in bind
    return self._true_bind(*args, **params)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/core.py", line 520, in _true_bind
    return self.bind_with_trace(prev_trace, args, params)
           ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/core.py", line 525, in bind_with_trace
    return trace.process_primitive(self, args, params)
           ~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/core.py", line 1024, in process_primitive
    return primitive.impl(*args, **params)
           ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/philipp/jax-metal/lib/python3.13/site-packages/jax/_src/dispatch.py", line 90, in apply_primitive
    outs = fun(*args)
jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: -:0:0: error: unknown attribute code: 22
-:0:0: note: in bytecode version 6 produced by: StableHLO_v0.19.0

--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
I0000 00:00:1741275406.517878 13032203 mps_client.h:209] MetalClient destroyed.

I tried different python versions as well as some other jax versions, but It won't run.

System info (python version, jaxlib version, accelerator, etc.)

jax: 0.5.2 (0.5.1 did not work either)
jaxlib: 0.5.1
numpy: 2.2.3
python: 3.12.8 | packaged by Anaconda, Inc. | (main, Dec 11 2024, 10:37:40) [Clang 14.0.6 ]
device info: Metal-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='Mac.fritz.box', release='24.3.0', version='Darwin Kernel Version 24.3.0: Thu Jan 2 20:24:23 PST 2025; root:xnu-11215.81.4~3/RELEASE_ARM64_T8122', machine='arm64')

[rajasekharporeddy]

Hi @PhilippFoerster

It seems jax-metal is not compatible with JAX 0.5.1 or later. As a temporary work around, installing jax-metal 0.1.1., jax 0.5.0 and jaxlib 0.5.0, resolves the issue. Thank you.

[jakevdp]

The last release of jax-metal was over a year ago now, and so it seems pretty clear that the team behind it is no longer working on it. We've removed mentions of jax-metal from JAX's installation docs: at this point I think it's probably best for users to consider the jax-metal project dead, and to focus on using supported backends.