JAX is NumPy with automatic differentiation, JIT compilation, and hardware acceleration baked in. It is not a deep learning framework in the way PyTorch is — there are no built-in layers or optimisers. Instead, JAX gives you the primitives for building them: pure functional transformations (jit, grad, vmap, pmap) that compose cleanly and run efficiently on GPU and TPU. This guide covers those primitives, the ecosystem libraries that build on them (Flax, Optax), and how to train your first neural network in JAX.
Installation
# CPU only
pip install jax
# NVIDIA GPU (CUDA 12)
pip install "jax[cuda12]"
# Apple Silicon (experimental MPS support)
pip install "jax[metal]"
# Verify
python -c "import jax; print(jax.__version__); print(jax.devices())"JAX Arrays and NumPy Compatibility
JAX arrays work like NumPy arrays but live on GPU by default and are immutable. Most NumPy operations exist in jax.numpy with the same API:
import jax
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
print(x.device()) # e.g. CudaDevice(id=0)
# NumPy-compatible operations
y = jnp.sin(x) + jnp.cos(x)
z = jnp.dot(x, x)
# Immutability — use functional updates
arr = jnp.zeros((3, 3))
arr = arr.at[0, 0].set(1.0) # not arr[0,0] = 1.0
print(arr)Automatic Differentiation: jax.grad
JAX differentiates pure Python functions with jax.grad(). It takes a function and returns a new function that computes the gradient:
import jax
def f(x):
return jnp.sum(x ** 2)
grad_f = jax.grad(f) # gradient function
print(grad_f(jnp.array(3.0))) # 6.0
# value_and_grad: compute f(x) and grad(f)(x) in one pass
val, grads = jax.value_and_grad(f)(jnp.array([1.0, 2.0, 3.0]))
# Higher-order gradients
hessian = jax.jacfwd(jax.jacrev(f)) # second derivative
# Gradient w.r.t. multiple arguments
def loss(params, x, y):
return jnp.mean((params['w'] @ x - y) ** 2)
grad_loss = jax.grad(loss, argnums=0) # grad w.r.t. params onlyJIT Compilation: jax.jit
Wrap any pure JAX function with jax.jit() to compile it for your hardware. The first call triggers compilation; subsequent calls run the compiled version:
import jax, jax.numpy as jnp, time
def slow_fn(x):
return jnp.sum(jnp.sin(x) ** 2 + jnp.cos(x) ** 2)
fast_fn = jax.jit(slow_fn)
x = jnp.ones(1_000_000)
# Warmup (compilation)
_ = fast_fn(x).block_until_ready()
t0 = time.time()
for _ in range(100):
_ = fast_fn(x).block_until_ready()
print(f"JIT: {(time.time()-t0)*10:.2f}ms per call")
t0 = time.time()
for _ in range(100):
_ = slow_fn(x).block_until_ready()
print(f"Eager: {(time.time()-t0)*10:.2f}ms per call")Vectorisation: jax.vmap
vmap vectorises a function over a batch dimension. Write your function for a single example; vmap handles the batching automatically — often more clearly than manual batching and as fast as a hand-written batched version:
import jax, jax.numpy as jnp
def predict_one(params, x):
"""Prediction for a single example."""
return jnp.dot(x, params['w']) + params['b']
# Vectorise over x (axis 0), keep params fixed (None)
predict_batch = jax.vmap(predict_one, in_axes=(None, 0))
params = {'w': jnp.ones(10), 'b': 0.0}
X_batch = jnp.ones((32, 10)) # 32 examples
preds = predict_batch(params, X_batch) # shape: (32,)
# Combine with grad for per-example gradients
per_example_grad = jax.vmap(jax.grad(loss_fn, argnums=0), in_axes=(None, 0, 0))Random Numbers in JAX
JAX handles random numbers differently from NumPy. Every operation that uses randomness requires an explicit key (PRNG key), and you must split keys to generate independent streams. This ensures reproducibility and compatibility with JAX’s functional model:
import jax, jax.numpy as jnp
key = jax.random.PRNGKey(42) # seed
# Always split before using
key, subkey = jax.random.split(key)
x = jax.random.normal(subkey, shape=(3, 4))
# Multiple random operations
key, k1, k2 = jax.random.split(key, 3)
a = jax.random.uniform(k1, shape=(5,))
b = jax.random.normal(k2, shape=(5,))Training a Neural Network with Flax and Optax
Raw JAX has no layers or optimisers — use Flax for the network and Optax for the optimiser:
import jax, jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state
class MLP(nn.Module):
@nn.compact
def __call__(self, x):
x = nn.relu(nn.Dense(256)(x))
x = nn.relu(nn.Dense(128)(x))
return nn.Dense(10)(x)
# Initialise
model = MLP()
key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, 784)))['params']
tx = optax.adam(1e-3)
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)
# Training step — compiled for speed
@jax.jit
def train_step(state, x, y):
def loss_fn(params):
logits = state.apply_fn({'params': params}, x)
loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
return loss, logits
(loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
state = state.apply_gradients(grads=grads)
return state, loss
# Training loop
for epoch in range(10):
for x_batch, y_batch in dataloader:
state, loss = train_step(state, jnp.array(x_batch), jnp.array(y_batch))
print(f"Epoch {epoch+1}: loss={loss:.4f}")Key Concepts to Internalise
JAX has four rules that trip up newcomers. First, functions passed to jit() must be pure — no side effects, no mutation of external state. Second, all arrays are immutable — use .at[].set() for updates. Third, every random operation needs an explicit key — split the key before each use, never reuse. Fourth, Python control flow inside JIT-compiled functions must be static — if conditions based on traced JAX arrays are not allowed; use jax.lax.cond() instead for dynamic branching inside JIT. These constraints feel restrictive at first but enable JAX’s compiler to reason about the computation graph completely, producing the aggressive optimisations that make it fast.
JAX rewards the effort of learning its functional model. Once the core transforms — grad, jit, vmap, pmap — become natural, you can express complex training algorithms in fewer lines than PyTorch, with the compiler handling the performance optimisation. The best way to start is with the official JAX tutorial at jax.readthedocs.io, which walks through these concepts interactively, followed by the Flax MNIST example which shows the full training loop in the style you will use for real projects.
Understanding JAX’s Functional Model
The key insight for JAX beginners: think of JAX as a compiler for mathematical functions, not an imperative programming environment. When you call jax.jit(f), JAX traces through f with abstract values to build a computation graph, then compiles that graph for your hardware. This means f must be a pure function — same inputs always produce same outputs, no mutation of state outside the function. The constraints are real but the payoff is total: once JAX has the computation graph, it can optimise it aggressively, vectorise it with vmap, differentiate through it with grad, and parallelise it with pmap, all composably.
Debugging JAX Code
Debugging is harder in JAX than PyTorch because JIT-compiled functions are traced, not executed eagerly. Common patterns: temporarily remove @jax.jit decorators to run in eager mode and inspect intermediate values; use jax.debug.print() inside JIT functions to print during execution; use jax.debug.breakpoint() to pause inside a compiled function. For shape errors, JAX’s error messages have improved significantly in recent versions and include the operation that caused the mismatch.
@jax.jit
def my_fn(x):
jax.debug.print("x shape: {}", x.shape) # prints during JIT execution
return jnp.sum(x)
# Or run without JIT for full Python debugging
result = my_fn.lower(x).compile()(x) # inspect compiled artifactMulti-Device Training with pmap
JAX’s pmap parallelises a function across multiple GPUs or TPUs — each device gets a shard of the data, runs the function in parallel, and collective operations (like all-reduce for gradient averaging) happen automatically:
import jax
n_devices = jax.device_count()
print(f"Devices: {jax.devices()}")
# Replicate params across devices
from flax.jax_utils import replicate, unreplicate
params = replicate(state.params)
# pmap over batch dimension
@jax.pmap
def parallel_train_step(params, x, y):
loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
grads = jax.lax.pmean(grads, axis_name='batch') # average gradients
return loss, grads
# Reshape batch: (n_devices, per_device_batch_size, ...)
x_sharded = x.reshape(n_devices, -1, *x.shape[1:])
y_sharded = y.reshape(n_devices, -1)JAX Ecosystem at a Glance
JAX itself provides only the core transformations and NumPy-compatible array operations. The surrounding ecosystem provides everything else. Flax is the most widely used neural network library — used by Google Research and increasingly by the broader JAX community. Optax provides gradient-based optimisers (Adam, SGD, schedule-based variants) with the same composability as JAX’s function transforms. Orbax handles checkpointing — saving and loading model state in a format compatible with JAX’s functional style. Equinox (from Patrick Kidger) offers a more PyTorch-like object-oriented feel built on top of JAX, which many practitioners find easier to learn than raw Flax. For probabilistic programming, NumPyro and BlackJAX provide Bayesian inference tools built entirely in JAX.
How jax.grad Works Under the Hood
Understanding how JAX differentiates functions helps avoid common mistakes. When you call jax.grad(f)(x), JAX traces f symbolically — it records every operation applied to x as a computation graph. Then it applies the chain rule backwards through that graph to compute df/dx. This is reverse-mode automatic differentiation, identical in principle to PyTorch’s autograd but implemented as a function transformation rather than as bookkeeping attached to tensor objects. The key difference: in PyTorch, gradient computation state is attached to tensors and accumulates in-place; in JAX, gradients are computed by transforming functions, with no persistent state. This makes JAX gradients composable — you can take the gradient of a gradient-computing function — but means you must explicitly pass parameters as function arguments rather than using module attributes.
When NOT to Use jit
Knowing when to avoid jax.jit is as important as knowing when to use it. Do not JIT functions with Python-side effects (print, file I/O, logging) — they will not execute during JIT-compiled runs. Do not JIT very small functions called only once — the compilation overhead outweighs the runtime benefit. Do not JIT functions with highly dynamic shapes that change every call — each new shape triggers recompilation. For iterative debugging — trying to understand why gradients are zero or why values are NaN — remove JIT temporarily, reproduce the issue in eager mode, fix it, then re-enable JIT. NaN debugging is much harder inside a JIT-compiled function because the computation happens asynchronously and error messages may be cryptic.
Saving and Loading Model State
JAX model state is just a Python dict of arrays (the parameter tree). Save it with Orbax or with simple pickle/numpy serialisation:
import orbax.checkpoint as ocp
import jax.numpy as jnp
# Save with Orbax
checkpointer = ocp.PyTreeCheckpointer()
checkpointer.save('/path/to/checkpoint', state.params)
# Restore
restored = checkpointer.restore('/path/to/checkpoint')
# Simple alternative with numpy
import numpy as np
np.save('params.npy', jax.tree_util.tree_map(np.array, state.params))
params_loaded = np.load('params.npy', allow_pickle=True).item()Common JAX Mistakes and How to Avoid Them
The mistakes that slow JAX beginners down most. Reusing PRNG keys — every call to a random function needs a fresh key split from the previous one. Reusing a key gives different functions the same random stream, producing silent correlations. Modifying arrays in-place — JAX arrays are immutable; arr[0] = 1 raises an error. Use arr.at[0].set(1) instead. Python control flow in JIT — if jax_array > 0 inside a JIT function will either always evaluate the same branch (traced with abstract values) or raise an error. Use jax.lax.cond() for dynamic branching inside compiled functions. Not calling block_until_ready() in benchmarks — JAX operations are asynchronous; timing without .block_until_ready() measures dispatch time, not computation time. Shapes must be static in JIT — functions that change array shapes based on data values (not just size) need either dynamic=True or to be restructured to use static shapes throughout.
JAX rewards the effort of learning its functional model. Once the core transforms — grad, jit, vmap, pmap — become natural, you can express complex training algorithms more concisely than in PyTorch, with the compiler handling performance. Start with the official JAX tutorial at jax.readthedocs.io, then the Flax MNIST example, then try reimplementing a training loop you already understand from PyTorch. That path — familiar task, new framework — is the fastest route to JAX fluency. The mistakes listed above are the ones every JAX beginner hits; internalising them upfront saves hours of debugging and makes the learning curve significantly less steep than it would otherwise be.
JAX is one of those frameworks that feels awkward for the first week and then clicks completely — the functional model is consistent and principled in a way that makes complex things like multi-device training and higher-order gradients feel natural rather than bolted on.