metaljax Brings Unmodified JAX Code to Apple Silicon GPUs

metaljax is a new open-source PJRT plugin that runs unmodified JAX code on Apple Silicon GPUs by interpreting compiled StableHLO programs onto MLX arrays. It behaves as a standard JAX backend, supporting jit, grad, vmap, lax.scan and optax training loops, with transformer training steps running within a few percent of PyTorch's MPS backend. Notably, the author built most of it through 'vibe coding'—AI-assisted iterative development—yet the project still reports passing 98.4% of the upstream JAX test suite (27,779 of 28,200 tests), with every release gated by a correctness sweep against the CPU backend.

This matters for engineers because JAX currently has no official Metal backend, leaving Apple GPU users dependent on PyTorch's MPS support. The known gaps are deliberate platform constraints rather than bugs: no float64 (Metal GPUs lack f64 ALUs, with opt-in f32 emulation), single-device-only pmap/shard_map, and denormals flushing to zero. Remaining issues—ordered-effect token threading, complex64 edge cases at poles, and a handful of PJRT surface APIs—are documented and individually audited against CPU parity.

The project is in beta, installable via pip, and defaults to CPU unless JAX_PLATFORMS=metal is set, requiring Apple silicon, macOS 14+, and Python 3.12+. The author notes the package will be deprecated if an official Metal backend ever lands upstream in JAX.