jax-metal

The jax-metal package is a Metal GPU plugin to provide Metal acceleration on Mac platforms for JAX applications.

Installation

The following table tracks jax-metal versions and compatible versions of jax, jaxlib and MacOS.

We recommend to install the binary package with venv or conda.

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

Usage

python -c 'import jax; print(jax.numpy.arange(10))'

Compatibility with jaxlib

jax-metal is compatible with the minimal jaxlib version tracked in the above table. It can be compatibly run with jaxlibs beyond the minimum version by setting the environment variable to ENABLE_PJRT_COMPATIBILITY=1.

pip install -U jaxlib jax
ENABLE_PJRT_COMPATIBILITY=1 python -c 'import jax; print(jax.numpy.arange(10))'

Currently not supported

The Metal plug-in is experimental and not all JAX functionality may be supported. Issues that are reported and tracked can be found in the list: https://github.com/google/jax/issues?q=is%3Aissue+is%3Aopen+metal

• Unsupported data types: np.float64, np.complex64, np.complex128
• The Metal plug-in doesn’t pass all tests under https://github.com/google/jax/tree/main/tests.

Please refer to https://developer.apple.com/metal/jax/ for the full setup and verification.