Why JAX Training Loops Are Explicit

In Keras, you call model.fit(x, y, epochs=10) and training happens magically. In PyTorch, you write a loop but the optimizer manages state internally. In JAX, everything is explicit — there's no model.fit(), no hidden optimizer state, no implicit backward pass. You write every step yourself.

This sounds tedious, but it's actually liberating. You have full control over:

• What gets differentiated and what doesn't
• How gradients are transformed before being applied
• When and how checkpoints are saved
• How data flows between devices
• Which parts of the model are frozen

Here's the fundamental JAX training pattern:

import jax
import jax.numpy as jnp

# The pattern: forward → loss → grad → update → repeat
def train_step(params, x, y, lr=0.01):
    # 1. Define loss as a function of params
    def loss_fn(params):
        predictions = model_forward(params, x)
        return jnp.mean((predictions - y) ** 2)

    # 2. Compute loss and gradients simultaneously
    loss, grads = jax.value_and_grad(loss_fn)(params)

    # 3. Update parameters (simple SGD)
    new_params = jax.tree.map(lambda p, g: p - lr * g, params, grads)

    return new_params, loss

# 4. JIT-compile for speed
train_step_jit = jax.jit(train_step)

# 5. Training loop
for epoch in range(num_epochs):
    for batch in data_loader:
        params, loss = train_step_jit(params, *batch)

That said, writing SGD by hand every time gets old fast. That's where Optax comes in — it handles the optimizer logic while keeping everything explicit and functional.