JAX against PyTorch

Eager tape autograd against traced, compiled pure functions. The same MNIST MLP in both, with APIs and gotchas.

Tasks and commands
Annotated code
Diagrams
Comparisons
Troubleshooting
Questions and tips

Side by side#

Factor	PyTorch	JAX
Autodiff	Tape-based autograd	Functional transforms
Execution	Eager (C++ kernels)	Traced + JIT compiled
Model	`nn.Module`	Pure function: `f(params, x) -> y`
TPU support	Community-maintained (torch_xla)	First-class
GPU support	Best-in-class (native CUDA)	Good (CUDA via XLA)
Debugging	Easy (eager, line-by-line)	Hard (tracing + compilation)
Speed (MNIST MLP)	~0.5s/epoch	~0.3s/epoch
Multi-device	`FSDP`	`jax.pmap(f)`

From Lesson 03.11, Lesson 03.12

Same task, two APIs#

Task	PyTorch	JAX
Get gradients	`loss.backward()`	`jax.grad(loss_fn)(params, x, y)`
Clear gradients	`optimizer.zero_grad()`	Not needed. Grads are return values
Update params	`optimizer.step()`	`optax.apply_updates(params, updates)`
Adam	`torch.optim.Adam(model.parameters(), lr=1e-3)`	`optax.adam(learning_rate=1e-3)`
Batch over examples	`for x in batch:`	`jax.vmap(f)`
Hold params	`model.parameters()`	Immutable pytree of arrays
Seed randomness	`torch.manual_seed`	`jax.random.PRNGKey(42)`
Save weights	`torch.save(model.state_dict(), "model.pt")`	`checkpointer.save('/tmp/model', params)`

From Lesson 03.11, Lesson 03.12

Two execution models#

From Lesson 03.11, Lesson 03.12

PyTorch: MNIST MLP loop#

model = MNISTModel().to(device)784, 256, 128, 10

criterion = nn.CrossEntropyLoss()takes raw logits

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):

model.train()dropout on

for images, labels in train_loader:

images, labels = images.to(device), labels.to(device)

optimizer.zero_grad()grads accumulate

outputs = model(images)

loss = criterion(outputs, labels)

loss.backward()replays the tape

optimizer.step()

torch.save(model.state_dict(), "mnist_mlp.pt")~97.8% test accuracy

From Lesson 03.11, 03.11 code/pytorch_intro.py

JAX transforms#

Task	Code
Differentiate a function	`jax.grad(f)`
Second derivative	`jax.grad(jax.grad(f))`
Batch a one-example function	`jax.vmap(predict, in_axes=(None, 0))`
Per-example gradients	`jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0))`
Split the batch over devices	`jax.pmap(train_step, axis_name='devices')`
Update every leaf	`jax.tree.map(lambda p, g: p - lr * g, params, grads)`
Set one element	`a = a.at[0].set(5)`
Clip, then Adam	`optax.chain(optax.clip_by_global_norm(1.0), optax.adam(1e-3))`

From Lesson 03.12

One JAX training step#

From Lesson 03.12, 03.12 code/jax_intro.py

JAX: MNIST MLP loop#

optimizer = optax.adam(learning_rate=1e-3)

@jax.jitcompile to XLA

def train_step(params, opt_state, x, y):

loss, grads = jax.value_and_grad(loss_fn)(params, x, y)

updates, opt_state = optimizer.update(grads, opt_state, params)

params = optax.apply_updates(params, updates)

return params, opt_state, lossnew state out

key = random.PRNGKey(0)

params = init_params(key)nested dict pytree

opt_state = optimizer.init(params)

for epoch in range(n_epochs):

key, subkey = random.split(key)new key each epoch

perm = random.permutation(subkey, len(X_train))

params, opt_state, loss = train_step(params, opt_state, xb, yb)batch_size 128

From 03.12 code/jax_intro.py, Lesson 03.12

Install and run#

Command	What it does
`pip install jax jaxlib optax flax`	Installs JAX, Optax and Flax
`pip install jax[cuda12]`	Installs JAX with GPU support
`pip install jax[tpu] -f <url>`	Installs JAX for TPU on Google Cloud
`XLA_PYTHON_CLIENT_PREALLOCATE=false`	Stops the 75% GPU memory preallocation
`jax.debug.print()`	Prints inside a jitted function
`fast_fn(x).block_until_ready()`	Waits for the result before you stop a timer
`torch.load("model.pt", weights_only=True)`	Loads a saved state dict

From Lesson 03.12, 03.12 code/jax_intro.py, Lesson 03.11

Gotchas#

Symptom	Cause	Fix
`Expected all tensors to be on the same device`	Model and batch on different devices	Move both with `.to(device)`
Wrong gradients and no error	Softmax output into `CrossEntropyLoss`	Pass raw logits
First step is slow	JIT compiles on the first call	Warm up before you benchmark
`if x > 0` fails inside JIT	Control flow depends on a traced value	Use `jax.lax.cond`
Python loop inside JIT is slow	Loop over JAX arrays	Use `jax.lax.scan` or `jax.lax.fori_loop`
`a[0] = 5` fails	JAX arrays are immutable	Write `a = a.at[0].set(5)`

From Lesson 03.11, Lesson 03.12

Use X when#

Use PyTorch unless you have a specific reason to use JAX.
Use PyTorch to debug line by line. print(), pdb and breakpoint() work.
Use PyTorch to prototype fast and for the HuggingFace ecosystem.
Use JAX when you have TPU access.
Use JAX for per-example gradients.
Use JAX for multi-device training at massive scale.

From Lesson 03.12, Lesson 03.11

Questions#

Why does JAX need pure functions?: JIT compilation requires them. Same inputs give the same outputs.
Why is there no zero_grad in JAX?: jax.grad returns grads as values. Tensors store no graph.
Why save state_dict(), not the model?: The model object uses pickle. Pickle breaks when you refactor.
Why torch.no_grad() in eval?: It disables autograd. Memory use goes down.
What maps from the mini framework?: Sequential, Adam, zero_grad, backward, step, train, eval.

From Lesson 03.10, Lesson 03.11, Lesson 03.12