jax-mac-gpu

Reproducible compatibility profiles and diagnostics for running JAX on Apple GPUs from a MacBook.

This project does not fork JAX or patch Metal. It keeps the backend choice in an isolated environment, reports the actual device, runs a small JIT/gradient check, and provides a synchronized CPU/MPS benchmark. A failed Metal check is reported; it is never silently replaced by CPU.

Why this exists

The Apple JAX Metal path is version-sensitive. A JAX installation can import successfully while still failing at StableHLO/PJRT initialization, compiling a specific operation, or running gradients. The goal of this repository is to make those boundaries easy to reproduce for other Mac users.

Profiles

Profile	Intended use	Status
metal	Current Apple Silicon/macOS profile using jax-metallib	Tested forward kernels on Apple M4; gradient support remains experimental
cpu	Reference and CI backend	Tested
legacy-metal	Older Apple jax-metal compatibility line	Reference only; not the default current profile

The current tested local profile is CPython 3.13.7, macOS 26.5.2, Apple M4, jax==0.10.1, jaxlib==0.10.1, and jax-metallib==0.10.1.0.

Install

Use Python 3.13 in a dedicated virtual environment:

python3.13 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip
python -m pip install -e ".[metal]"

For CPU-only use:

python -m pip install -e ".[cpu]"

The legacy profile is opt-in and should be tested against the target macOS/JAX combination before use:

python -m pip install -e ".[legacy-metal]"

Diagnose

jax-mac-gpu doctor --backend mps
jax-mac-gpu doctor --backend mps --strict
jax-mac-gpu doctor --backend cpu --strict

doctor distinguishes device discovery, forward JIT, and gradient execution. For experimental Metal releases, a partial result is useful evidence: it means the device or forward path works while another operation is unsupported or unstable. Use --strict in CI or before enabling a training workflow.

Benchmark

jax-mac-gpu benchmark --backend cpu --iterations 10
jax-mac-gpu benchmark --backend mps --iterations 10

Every timed call uses block_until_ready(). The benchmark is intentionally small and forward-only; it is not evidence that a full training job or every JAX primitive works on Metal. Compare CPU and MPS for the actual workload.

License and dependencies

The wrapper code is MIT-licensed. JAX, jax-metal, and jax-metallib remain separate dependencies with their own licenses and limitations. Apple’s JAX Metal documentation describes the plugin as experimental; see the upstream references before selecting a profile:

• Apple accelerated JAX on Mac
• jax-metal on PyPI
• jax-metallib on PyPI