metaljax — a Metal backend for JAX

metaljax lets unmodified JAX code run on Apple-silicon GPUs:

$ JAX_PLATFORMS=metal python -c \
    "import jax.numpy as jnp; a = jnp.array([1, 2, 3]); print(2 * a, (2*a).device)"
[2 4 6] MetalDevice(id=0)

From JAX's point of view it is a regular PJRT backend: jax.devices() reports a METAL device, and jit, grad, vmap, lax.scan, jax.random (threefry), optax training loops, etc. all work. Under the hood a native plugin lowers the compiled StableHLO programs onto MLX arrays, which execute on the GPU via Metal. The wheel carries its own patched MLX runtime, so nothing else needs installing.

Status: beta. Twenty real models — LLM decode (dense, MoE, quantized), a vision encoder, diffusion, LoRA and full-parameter training — run end-to-end through unmodified JAX code, and are re-measured on the binary of every release (the table under Benchmarks, the ledgers models.md and STATUS.md). Training steps beat PyTorch's MPS backend on both training rows; decode trails the Metal-native inference stacks (mlx-lm, llama.cpp) by 1.1–1.9×. Every release is gated by the pinned JAX test suite (jax 0.11.2, 99.53 % passing), a whole-model correctness sweep against the CPU backend (106/106), and the model battery. Coverage gaps remain — unsupported constructs are declined at compile time, naming the op. If a Metal backend ever lands upstream in the JAX ecosystem, this package will be deprecated in its favor.

What's new in 0.11.9

• Training loops no longer hold many steps' worth of memory at once. A counted loop (a lax.scan over training steps) used to decide how much work to queue ahead of the GPU by op count alone, and MLX allocates every intermediate when work is queued — so a 10k-weight recurrent LM's 256-step training chunk briefly held 5 GB for a 0.3 GB working set and died with RESOURCE_EXHAUSTED on a busy 32 GB machine. Loops now measure what each queued step holds and keep at most METALJAX_LOOP_INFLIGHT_MB (1 GB) in flight, waiting on the oldest step so the GPU stays busy: that chunk now peaks at 1.1 GB and runs ~17 % faster (the old cadence drained the GPU at every sync), and texmo's own multi-model training loop peaks at 4.9 GB instead of 8.5.
• The memory governor gives back its own cache before it gives up. Past a hard line it now waits for in-flight GPU work to finish and clears MLX's buffer cache, at the start of the stall and on every tick of it, instead of once and rate-limited; a refusal now says how much of what the process holds is live arrays and how much is cache. Details: notes/loop-inflight-window-2026-09.md.

What's new in 0.11.8

• jax 0.11.2. The pinned release (and its test suite) moved from jax 0.11.0 to 0.11.2; the plugin keeps working with 0.11.0/0.11.1. Moving surfaced a silent-wrongness bug fixed in this release: floating-point lax.rem / jnp.fmod was inexact on every earlier release (up to 1 ULP of the dividend in f32, large errors in bf16, and overflow to inf in f16); it is now bit-exact against the CPU backend. Also new: int2/uint2 (and sub-byte integer conversions now saturate like jax-CPU), and with_layout_constraint is accepted. Note jax 0.11.2 itself refuses bf16/f16 linear algebra at trace time. The plugin now builds against jax 0.11.2's own XLA commit (PJRT C API 0.115, StableHLO/VHLO 1.20, so jaxlib no longer downgrades programs to an older VHLO for it); jaxlib 0.11.0 still loads it. Details: notes/jax-0.11.2-pin.md.
• Integer shifts, popcnt and clz match XLA bit for bit (a silent-wrongness fix): lax.shift_right_arithmetic on unsigned integers ran as a logical shift, negative or very large shift amounts escaped the saturation guard into Metal's mod-width shift, and the emulated 2/4-bit integers shifted and counted their 8-bit storage. All are now exact against the CPU backend.
• Precision default. METALJAX_MATMUL_PRECISION=high is now the shipped default: f32 GEMM/attention stay off the M5's neural accelerators (exact, the pinned ~1e-6-vs-f64 class), while bf16/f16 and quantized matmuls use them — what mlx-lm and torch-MPS do. highest pins the pre-accelerator arch for every dtype; default is MLX's own (accelerator f32, ~8e-4). f32 is never silently degraded.
• Decode and MoE work. A ragged decode form for the MoE expert dispatch (a direct gather_mm over sorted expert ids); stacked sibling packs (several stacked per-layer dots reading one activation become one gather_mm over their [L, K, N] views concatenated on the free axis); projection packs (sibling decode projections over one activation become one dot over concatenated weights); the rotate-half rope as one fused kernel over a reversed pair view; submit-ahead pipelining of dynamic while loops; the KV cache written in place; chunk boundaries donating their carries; loop-position specialization (one compiled body per chunk with the loop position folded into its dynamic slice starts, now on by default). In the vendored MLX fork: empty command buffers skipped in gpu::finalize, and a transient-only byte cadence (opt-in, off by default).
• Recognizer coverage: a fused gated-delta-net decode step; the norm recognizer covers keras-hub's Qwen3.5/gemma4 and flax-NNX spellings; fused attention on the keras MoE rows (sinks, chained masks) and on multi-span GQA decode.
• keras-hub's Gemma4 int4 models are numerically broken upstream. The Gemma4 decoder block's "HOTFIX" multiplies activations by the raw unpacked int4 codes, scale unapplied, so every backend decoded one repeated token; every row-13 cell through 0.11.7 was timing that broken graph. The benchmark harness now routes the int4 FFN through the quantized layers and gathers embedding rows before unpacking (scripts/model_bench/int4_fix.py). Write-up: notes/keras-hub-gemma4-int4-hotfix-issue.md.
• Two msl_scan silent-wrongness fixes: a value that is one scalar per lane is not a register vector (it was broadcast over a whole row of the buffer), and only the proven induction variable is a loop counter — every other scalar carry is now classified by what the body does with it (invariant carries pass through; stepped ones become kernel state). Both are described in CLAUDE.md and in their commits; no shipped number or token stream moved.
• New comparator targets in the ledgers: mlx-lm on a local bf16 Mixtral 8×7B conversion (row 12), mlx-lm's own 4-bit quantization of gemma4-E2B (row 13), and a torch-MPS full AdamW train step at MaxText's defaults (row 19).

Install

pip install metaljax

Requirements: Apple-silicon Mac, macOS 14+, Python 3.12+, jax 0.11.x (installed automatically). No mlx install is needed — the wheel bundles its own — and if you do have the public mlx, the two coexist. Then select the backend per program:

JAX_PLATFORMS=metal python -c "import jax; print(jax.devices())"

CPU remains the default backend when JAX_PLATFORMS is unset, so installing metaljax does not change existing workflows. metaljax ships as a wheel only: the plugin is built against a pinned XLA workspace with bazel, which is not something an sdist can compile at install time.

How it works

jax.jit(f)(x)
  │  StableHLO (serialized portable artifact)
  ▼
libmetal_pjrt_native.dylib    ── the plugin: an xla::PjRtClient, with
  │  plugin-native/metal/        XLA's pjrt_c_api_wrapper_impl making the
  │                              PJRT C API around it. No Python, no GIL.
  │  compile: parse + lower StableHLO to a tape, decide what to fuse
  │  execute: replay the tape on device buffers
  ▼
plugin-native/runtime/        ── the executor: op emitters, MSL kernel
  │                              codegen, control flow, host LAPACK
  ▼
libmlx_metaljax.dylib         ── our vendored, privately install-named
                                 MLX: lazy Metal arrays, unified memory

• The plugin is self-contained: LLVM/MLIR/StableHLO/absl are linked in and private, and exactly two symbols are exported (GetPjrtApi and a callback bridge), so it coexists with TensorFlow-class carriers in one process.
• Registration happens through the jax_plugins namespace package (src/jax_plugins/metal/), at priority −1: CPU stays the default backend unless you opt in via JAX_PLATFORMS. That module is also where jax.debug.print / pure_callback callables live: the plugin calls back into them through one C function pointer.
• Python appears nowhere on the execute path.
• Before the first program runs, the plugin lowers what it can into fused forms: quantized matmuls, mixture-of-expert dispatches, softmax attention, norms, rope, sibling projection and per-layer stack packs, and — for recurrent scan bodies — generated Metal kernels. Each recognizer has an off switch (see Environment variables), and each is checked against the literal graph by the differential suites.

Requirements

• Apple-silicon Mac (developed on an M5 Max, macOS 26.5, Xcode 26.6 — any arm64 Mac with a recent Xcode/CLT should work).
• uv (only for creating the venv).
• Python 3.14 and jax/jaxlib 0.11.x (tested against 0.11.2, the pinned release, whose XLA commit the plugin is built against; jaxlib 0.11.0 still works — it negotiates the newer plugin down).

Developing from source

git clone https://github.com/eterevsky/metaljax && cd metaljax
uv venv --python 3.14 .venv
uv pip install -p .venv/bin/python jax numpy pytest
uv pip install -p .venv/bin/python -e .
./scripts/vendor_mlx.sh                        # build + stage the MLX runtime
cd plugin-native && bazel build //metal:libmetal_pjrt_native.dylib && cd ..
cp plugin-native/bazel-bin/metal/libmetal_pjrt_native.dylib src/metaljax/lib/

(The first bazel build is ~7 minutes; after that it is seconds. An editable install loads the dylib out of src/metaljax/lib/, and METALJAX_PLUGIN_PATH overrides it with any build you want to measure.)

Verify:

JAX_PLATFORMS=metal .venv/bin/python -c "import jax; print(jax.devices())"

should print [MetalDevice(id=0)].

Running the tests

The pytest suite runs everything through the real plugin — jax.jit on the Metal device, or compile_and_load for hand-written StableHLO — and compares against the JAX CPU backend:

.venv/bin/python -m pytest tests/ -q

Current suite: 512 tests (+1 xfail) across elementwise/transcendental ops, shapes and broadcasting, dot_general/einsum, reductions and cumulative ops, control flow (while/cond/scan), gather/scatter, sorting, convolutions, linalg, complex, RNG, sub-byte and bf16/f16 dtypes, the quantized-matmul / MoE / attention recognizers, the generated scan kernels, donation, buffer pointers, concurrency and the Metal command-buffer canaries.

The plugin has its own differential suites, which compare it against jax-CPU expression by expression and on whole models:

.venv/bin/python plugin-native/execute_test.py     # vs jax-CPU
.venv/bin/python plugin-native/texmo_gate.py       # whole-model training
cd plugin-native && bazel test //...               # C++ unit tests

End-to-end smoke test (device buffers, compile, execute, PJRT events):

JAX_PLATFORMS=metal .venv/bin/python -c "
import jax, jax.numpy as jnp
g = jax.jit(jax.grad(lambda x: jnp.sum(jnp.tanh(x) ** 2)))(jnp.arange(4.0))
print(g, g.device)"

Coverage and known gaps

Running the test suite of the exact jax release we pin (v0.11.2) executes 28,814 tests with 99.53 % passing — 28,679 passed / 135 failed (plus 6,120 skipped), measured 2026-09-23 on the release binary. Against the previous release's set (on jax 0.11.0) nothing regressed: the 11 added ids are 9 complex-plane accuracy tests that jax-CPU fails identically, one f64 test and one TPU-only layout test, and the two genuine new failures jax 0.11.2 exposed were fixed (notes/jax-0.11.2-pin.md). It concentrates in export_harnesses_multi_platform_test (44), lobpcg_test (27), x64_context_test (13 — the f64 policy below), api_test (5), export_test (5), shape_poly_test (4), xla_transform_test (4) and async_collectives_test (3). Every remaining failure has been individually examined and classified with evidence (notes/jax-test-suite-2026-07.md); they fall into three groups:

Intentional (platform constraints, will not change):

• No float64. Metal GPUs have no f64 ALUs. f64 values may pass through the device (stored as f32), but f64 compute fails at compile time naming the op; METALJAX_F64=downcast opts into f32 emulation. Keep jax_enable_x64 off. Same policy for complex128.
• One physical device. pmap/shard_map/collectives work on a single device (replica groups of size 1); actual multi-device sharding has no hardware to run on.
• Denormals flush to zero on the GPU (hardware behavior); tests asserting subnormal outputs (e.g. jnp.spacing) differ from CPU.
• Complex special values at inf/NaN poles for a handful of transcendentals (log/trig/hyperbolic family) follow MLX's kernel semantics rather than C99. Finite inputs match CPU; full C99 pole behavior would need per-element branches in hot paths (policy: not worth the slowdown). sqrt/rsqrt/exp/expm1/tan/abs/sign are rebuilt and exact.

Remaining audited gaps (every one re-examined during the 0.11.0 parity campaign and approved as-is; each carries evidence in notes/jax-test-suite-2026-07.md):

• Ordered-effect residue: buffer_callback and emit_python_callback are rejected by jax-side platform allowlists (callback.py, buffer_callback.py hard-code cpu/cuda/rocm/tpu) — not reachable from a plugin; verified passing on CPU because cpu is inside those hard-coded lists. Ordered debug.print/io_callback work.
• testSincInfinities, FD-reference gradient corners: fail on the CPU backend too, or the test's finite-difference reference is numerically meaningless in f32 (documented with numbers).
• Better-than-reference cases: shape-polymorphic jnp.insert / jnp.nonzero — the harness asserts NotImplementedError because jax's CPU path cannot lower them; ours can, and values match CPU on concrete shapes. We fail these tests by succeeding.
• test_dce_sink_prevents_xla_dce: needs optimized-HLO text retrieval (PJRT_Executable_OptimizedProgram), a debugging surface we have not implemented.

Supported (each verified against the CPU backend): sorting (sort/argsort/top_k/approx_top_k/median/percentile/ unique, key-value and multi-key lexicographic sorts — jnp.lexsort, unique(axis=), set operations — IEEE total-order NaN handling, complex lexicographic order); convolutions (1/2/3-D float, integer — exact, and complex; strided, dilated, grouped, transposed, and their gradients); the full scatter family (windowed, out-of-bounds-dropping, arbitrary elementwise bodies); general reduce/reduce_window bodies and pooling with gradients (select_and_scatter, select_and_gather_add); complex64 end-to-end (arithmetic, FFT, linalg); linear algebra via LAPACK semantics on the host (QR, eigh, eig, SVD, LU, Cholesky, triangular_solve, Schur, Hessenberg — CPU-bound in every backend, free on unified memory) — including bfloat16/float16 inputs, which jax's CPU backend itself rejects (computed in f32, results in the requested dtype); single-device pmap/shard_map with the full collective set; rng_bit_generator (Philox and ThreeFry, bit-exact vs CPU, so the rbg/unsafe_rbg PRNG implementations work); int4/uint4 and all float8 dtypes (emulated: exact values in wider storage, grid-quantized converts, 4-bit wraparound); host callbacks (jax.debug.print, pure_callback, io_callback); shape-polymorphic jax.export of all of the above; popcnt/count_leading_zeros; sparse (BCOO/BCSR) workloads.

Unsupported constructs fail loudly at compile time with the op named — nothing silently falls back to CPU or returns wrong dtypes.

Using metaljax from another project

Add metaljax to your dependencies (it declares jax itself):

[project]
dependencies = ["metaljax"]

and set JAX_PLATFORMS=metal (or jax.config.update("jax_platforms", "metal") before first use).

To develop against a local checkout instead, use a path source:

[tool.uv.sources]
metaljax = { path = "../metaljax", editable = true }

(with an editable install, build the plugin once in the checkout — see Developing from source above).

Environment variables

The knobs worth knowing. Every flag the code reads — including the recognizer A/B switches and the debug-only ones — is listed with its category in notes/env-flags.md.

Repository layout

CLAUDE.md                  project decisions/status (kept current)
models.md                  the 21-row model ledger (per-release columns)
STATUS.md                  current model cells + cross-framework comparators
pyproject.toml             python package + jax_plugins entry point
plugin-native/             THE ENGINE (bazel workspace)
  metal/                   the PJRT plugin: xla::PjRtClient, StableHLO
                           ingest, lowering, the recognizers and packs
  runtime/                 the executor: op emitters, MSL kernel codegen,
                           control flow, host LAPACK, memory governor
  third_party/mlx/         our vendored MLX, linked privately
  execute_test.py          differential suite vs jax-CPU
  texmo_gate.py            whole-model training gate vs jax-CPU
src/jax_plugins/metal/     backend registration (priority -1) + the
                           host-callback registry the plugin calls into
src/metaljax/              __version__, and lib/ where the plugin dylib
                           and the vendored MLX runtime land
tests/                     pytest suite (Metal vs CPU, through PJRT)
notes/                     investigation notes, env-flags, gate records
scripts/                   benchmark, gate & release drivers
  model_bench/             the 21-row model battery + its manifest
  vendor_mlx.sh            build + stage the vendored MLX runtime
  build_native_wheel.sh    build + verify the release wheel

Benchmarks

Real models

Twenty models through unmodified JAX code, all measured 2026-09-27/28 on the 0.11.9 release binary (frozen-0119main-c8577cab, jax 0.11.2), one GPU process at a time with a cool-down between rows, timed through np.asarray (jax.block_until_ready is a no-op on this backend). goal is the best non-metaljax cell for that row at the same precision and workload — same checkpoint, prompt window, generated token count, batch, resolution and encoder set; custom kernels are fair game, different quantization is not. Provenance, per-row protocols, memory footprints, jax-CPU cells and every caveat live in STATUS.md and models.md.

Metric: LLM rows = warm decode ms/token; vision = forward ms; diffusion = ms/diffusion-step; training = ms/step. Lower is better; ratio < 1 means metaljax is ahead.

Reading the table: metaljax is ahead of PyTorch-MPS on every training and diffusion row, and behind the dedicated Metal inference stacks on decode — mlx-lm runs on the same Metal library underneath, so that band is the optimization target, and llama.cpp's hand-written kernels lead even mlx-lm on bf16. Rows without a goal cell have no like-for-like non-metaljax implementation (rows 14/15 are qwix-quantized JAX models). Several rows do not run on the JAX CPU backend at all at these sizes; where they do, the CPU cells are in STATUS.md (e.g. row 19: 1398 vs 362.6 ms/step, row 16: 363.6 vs 41.7 ms, row 2: 312.1 vs 58.1 ms/token). Row 20 (Qwen3-235B-A22B 3-bit, 45.2 ms/tok at 0.11.8) left the suite at 0.11.9: it needs ~104 GB of a 128 GB machine, more than a desktop session leaves free (STATUS.md fn 21).

Correctness for these rows is gated the same night: greedy token streams are compared against the jax-CPU backend, and every divergence is accounted for (at this gate: one certified-benign 1-bf16-ULP logit tie on gemma4-E2B, the accepted accumulation-order class; all other CPU-comparable rows agree exactly over 64 tokens).

Training and recurrent workloads

The acceptance workload is a 106-config language-model training suite (dense, GRU/LSTM-family and linear-RNN cells, from tens of weights to several million) plus a 223-config performance sweep. At this release: 106/106 correct — one jitted training chunk per config executed on both backends from identical inputs, every output leaf compared against jax-CPU at a 1-ULP sensitivity-scaled tolerance — and the perf sweep is 1.14× faster than the standing anchor over 223 matched configs (134 configs improved >5 %, one regressed >5 %).

How it gets there: pure programs and counted-loop (scan/fori_loop) bodies are traced once into a fused Metal graph and replayed; small statically-counted loops are unrolled into the enclosing trace, so a whole recurrent-model training step (forward scan + backward + AdamW) becomes a single graph replay. On top of that, recurrent scan bodies that pattern-match as elementwise/matvec cells (rnn/gru/mgru/lrnn/rglru family — forward and the AD-generated backward loop) compile to a single generated persistent Metal kernel: the whole scan is one kernel launch, with state in registers (small cells), register-block lanes (small block matvecs, in-lane reductions and narrow rectangular readouts), or one threadgroup per batch element with the feature dim as the thread axis (full-width cells like gru.256, including rectangular fused-gate dots). Very wide cells (gru.1024-class) deliberately stay on the compiled-graph path, where batched matmul wins. Weight-gradient accumulations are handled by loop fission: the kernel stacks per-step operands and the einsum runs as one batched matmul after it.

openxla/xla benchmark suite

The single-device benchmarks from xla/tools/benchmarks (HLO converted to StableHLO with xla-translate, run via scripts/run_stablehlo_bench.py; ms per call, identical seeded inputs, outputs cross-checked against the CPU results). These cells were measured on the 0.11.6 binary and have not been re-run since — they are kept for the cross-device comparison, not as current-release numbers:

¹ exceeds the 4090's 24 GB VRAM; the M5's 128 GB unified memory runs gemma3_12b (23.5 GB of bf16 weights) where the discrete GPU cannot. ² the imported modules contain one plain stablehlo.dot (the logits matmul), which the native plugin declines by design (jax never emits it); measured with that one op rewritten to the equivalent dot_general, validated end-to-end against CPU references from the pristine modules. ³ VACUOUS under the suite's seeded inputs: both are generate-loop programs whose while-loop runs zero iterations, so no forward pass executes on any backend — the cells measure loop-condition + state-copy overhead only. Kept for completeness. ⁴ compiled whole-graph; ~9× CPU.

Correctness vs CPU on identical inputs: the gemma3 family diverges ≤3.8 % in bf16 KV-cache tensors (a few bf16 ULPs across 26+ layers) — the 4090 shows the same divergence class vs CPU (≤4.2 %), so that's cross-backend bf16 numerics, not a backend bug. maxtext NaN placement matches CPU exactly on all 11 NaN-carrying outputs.

Known limitations

Three platform constraints are permanent (detailed under Coverage and known gaps above): no float64 or complex128 compute (pass-through is fine; METALJAX_F64=downcast emulates in f32), one physical device (single-device pmap/shard_map/collectives work; real multi-device sharding has no hardware), and denormals flushing to zero on the GPU.

Performance and operational, not correctness:

• Decode trails the dedicated Metal inference stacks by 1.1–1.9× on the model table above; prefill and model load trail by more (see STATUS.md's gap band). Decode is the optimized path.
• Large checkpoint restores can need the governor's machine line raised. The restore transient of the biggest rows sits above the shipped METALJAX_MEM_SYS_MB default — row 10 (DeepSeek-V2-Lite) runs at METALJAX_MEM_SYS_MB=107520 on a 128 GB machine. Under the default the governor refuses cleanly (RESOURCE_EXHAUSTED) rather than wedging the machine, which is the contract, but the run does not start.
• A big quantized model pays a one-time pack wave. Row 20 (Qwen3-235B-A22B 3-bit) spends ~58 minutes packing quantized weights at a ~103 GB peak before steady decode at 43.3 ms/token. It is once per process, not per step.
• keras-hub's Gemma4 int4/int8 quantization is broken upstream (still present as of 2026-09-19): the decoder block's "HOTFIX" multiplies by unscaled codes, producing a numerically broken model on every backend. Our benchmark harness works around it (scripts/model_bench/int4_fix.py); a user's own quantized Gemma4 will be wrong until keras-hub fixes it.
• Scan bodies that don't fit the kernel-codegen patterns (gather/ scatter in the loop, non-affine indexing, bodies exceeding the trace or binding budgets) fall back to per-timestep compiled-graph replay, which pays per-step dispatch.
• Buffer donation is honoured (donate_argnums invalidates the donated inputs, matching other backends), but MLX cannot write outputs into the donated memory in place — the win is prompt buffer release rather than CUDA-style aliasing.

License / provenance

Experimental personal project; builds against public JAX/OpenXLA (PJRT header vendored from openxla/xla) and Apple's MLX. Apache-2.0.