applejax

A JAX backend for Apple Metal Performance Shaders (MPS), enabling GPU-accelerated JAX computations on Apple Silicon.

Fork of tillahoffmann/jax-mps with full linear algebra, complex number support, comprehensive scatter/gather handling, and 2000+ tests.

Quick Start

pip install applejax

The plugin registers itself with JAX automatically. Set JAX_PLATFORMS=mps to select it explicitly.

Requires macOS on Apple Silicon, Python >= 3.11, and jaxlib 0.9.x.

Performance

applejax achieves a modest 3x speed-up over the CPU backend when training a simple ResNet18 model on CIFAR-10 using an M4 MacBook Air.

$ JAX_PLATFORMS=cpu uv run examples/resnet/main.py --steps=30
Time per step (second half): 3.041

$ JAX_PLATFORMS=mps uv run examples/resnet/main.py --steps=30
Time per step (second half): 0.991

What Works

All JAX operations are supported, verified across 2000+ tests covering all major categories:

Linear Algebra

All operations work for both real (float32) and complex (complex64) inputs:

ML Framework Compatibility

Tested successfully with:

• Flax NNX — training loops with optimizers
• NumPyro — MCMC inference (NUTS sampler)
• Optax — all standard optimizers
• Equinox — neural network modules

Transformer components work end-to-end: multi-head attention, RoPE, RMSNorm, SwiGLU, causal masking.

Known Limitations

These are Metal/MPS hardware constraints, not bugs in applejax:

Architecture

This project implements a PJRT plugin to offload evaluation of JAX expressions to a Metal Performance Shaders Graph. The evaluation proceeds in several stages:

1. JAX lowers the program to StableHLO, a set of high-level operations for ML.
2. The plugin parses the StableHLO representation and builds the corresponding MPS graph. The graph is cached to avoid re-construction on repeated invocations.
3. The MPS graph is executed on the GPU. Operations not natively supported by MPS (e.g., linear algebra decompositions) run on CPU via Apple's Accelerate framework using a "native handler" mechanism.

Operation Implementations

Building from Source

1. Install build tools and build LLVM/MLIR & StableHLO (one-time, ~30 minutes):

brew install cmake ninja
./scripts/setup_deps.sh

2. Build and install:

uv pip install -e .

3. Install dev dependencies (test runner, linters, ML frameworks used in examples) and run tests:

uv sync --all-groups
uv run pytest

Version Pinning

applejax is built against the StableHLO bytecode format matching jaxlib 0.9.x. The setup_deps.sh script pins LLVM and StableHLO to specific commits from the XLA version used by jaxlib 0.9.0.

Runtime compatibility: Any jaxlib 0.9.x release should work with a built binary — the bytecode format is stable within a minor version series. The plugin will warn (not error) if the minor version doesn't match.

Updating for a new jaxlib release: Trace the dependency chain:

# 1. Find XLA commit used by jaxlib
curl -s https://raw.githubusercontent.com/jax-ml/jax/jax-v0.9.0/third_party/xla/revision.bzl

# 2. Find LLVM and StableHLO commits used by that XLA version
curl -s https://raw.githubusercontent.com/openxla/xla/<XLA_COMMIT>/third_party/llvm/workspace.bzl
curl -s https://raw.githubusercontent.com/openxla/xla/<XLA_COMMIT>/third_party/stablehlo/workspace.bzl

Then update STABLEHLO_COMMIT and LLVM_COMMIT_OVERRIDE in setup_deps.sh.

Project Structure

applejax/
├── CMakeLists.txt
├── src/
│   ├── jax_plugins/mps/         # Python plugin: registration + lowering rules
│   ├── pjrt_plugin/             # C++ PJRT implementation
│   │   ├── pjrt_api.cc          # PJRT C API entry point
│   │   ├── mps_client.h/mm      # Metal client management
│   │   ├── mps_executable.h/mm  # StableHLO compilation & execution
│   │   └── ops/                 # Operation implementations
│   │       ├── unary_ops.mm     # Element-wise unary operations
│   │       ├── binary_ops.mm    # Binary operations, dot products, matmul
│   │       ├── bitwise_ops.mm   # Bitwise operations
│   │       ├── shape_ops.mm     # Gather, scatter, reshape, pad, etc.
│   │       ├── reduction_ops.mm # Reduce, reduce_window, scan
│   │       ├── linalg_ops.mm    # Cholesky, QR, SVD, eigh, eig, Schur
│   │       ├── convolution_ops.mm # Convolution
│   │       ├── control_flow_ops.mm # cond, while, scan
│   │       ├── fft_ops.mm       # FFT operations
│   │       ├── sort_ops.mm      # Sort and top-k
│   │       ├── tensor_creation_ops.mm # Constants, iota
│   │       └── registry.h       # Op registration macros
│   └── proto/                   # Protobuf definitions
├── tests/
│   ├── test_ops.py              # Main test file (parameterized)
│   └── configs/                 # Test configurations by category
└── mps_ops/                     # Reference docs for MPS Graph methods

Benchmarks

uv run pytest -m benchmark --benchmark-only