A full training loop in JAX

Setup#

pip install flax optax

These pull in JAX. Outputs verified with jax 0.6.2, flax 0.10.7, optax 0.2.8, CPU. Seeded, so this run reproduces; across versions the exact losses may drift.

The whole thing#

XOR as the dataset: four rows, unsolvable by any straight line, so the network must genuinely learn — and small enough that every number can be watched.

import jax
import jax.numpy as jnp
import flax.linen as nn
import optax


class TinyNet(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(16)(x)
        x = nn.relu(x)
        return nn.Dense(1)(x)


# XOR: the smallest dataset a linear model cannot solve
x = jnp.array([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = jnp.array([[0.], [1.], [1.], [0.]])

model = TinyNet()
params = model.init(jax.random.key(0), x)
optimizer = optax.adam(0.05)
opt_state = optimizer.init(params)


def loss_fn(params, x, y):
    logits = model.apply(params, x)
    return jnp.mean(optax.sigmoid_binary_cross_entropy(logits, y))


@jax.jit
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, loss


for step in range(401):
    params, opt_state, loss = train_step(params, opt_state, x, y)
    if step % 100 == 0:
        print(f"step {step:3d}  loss {loss:.4f}")

preds = jax.nn.sigmoid(model.apply(params, x))
print("predictions:", jnp.round(preds.ravel(), 2))

step   0  loss 0.6802
step 100  loss 0.0013
step 200  loss 0.0005
step 300  loss 0.0003
step 400  loss 0.0002
predictions: [0. 1. 1. 0.]

XOR: learned. Every lesson of this section is in the file — find them.

The walkthrough#

train_step is the whole story. State in — (params, opt_state) — state out, plus the loss for logging. It is a pure function: call it twice with the same inputs and get identical outputs. That purity is what lets @jax.jit compile the entire step — forward pass, backward pass, and Adam's arithmetic fuse into one program. The first call traces and compiles; the remaining 400 replay.

value_and_grad delivers the loss (for the log) and gradients (for the update) from one shared pass — the training-loop reason that function exists.

The model outputs logits — raw scores, no sigmoid — and optax.sigmoid_binary_cross_entropy folds the sigmoid into the loss, the numerically safe arrangement. Sigmoid gets applied explicitly only at prediction time, at the bottom.

The Python loop stays outside jit — fine, because each iteration is one compiled call, and the print is honest Python. Squeezing further (the loop itself as lax.scan, zero Python per epoch) is the standard next optimisation for big models; at this size it would buy nothing.

Reassignment is the baton pass. params, opt_state, loss = train_step(...) — nothing mutates; names rebind to new trees. Forget to catch the returns and training silently reprocesses the same initial state 400 times.

Scaling this skeleton up is remarkably direct: real projects add batching over a dataset, a key split per step for dropout, checkpointing of (params, opt_state) — but the five lines inside train_step survive almost unchanged into serious codebases (often bundled into a TrainState object from flax.training).

Common mistakes#

Dropping a returned baton. Catching params but not opt_state restarts Adam's memory every step — training crawls and nobody crashes. Audit the tuple.

Jitting a step that closes over changing values. train_step reads model and optimizer from the enclosing scope — safe, because both are fixed structure. Closing over something you change between calls (a Python counter, a flag) bakes the trace-time value in forever. Changing things enter as arguments.

Logging every step. loss is a device array; printing it forces a wait for the computation (async dispatch). Printing every step can serialise the pipeline — log every N steps, as the loop above does.

Recompiling per shape. Feed batches of varying size and each new shape retraces the step. Fixed batch shapes, pad the last batch — the jit rules never stop applying.

Try it yourself#

Break it on purpose, twice. First: swap optax.adam(0.05) for optax.sgd(0.05) and watch step-400 loss — momentumless descent on a non-convex problem. Second: restore Adam but change the init key to jax.random.key(7) — same architecture, different starting point, different (probably still successful) path. You now know which knob did what.

What to learn next#

• Backpropagation — the algorithm inside value_and_grad, in full.
• Building models with Flax — revisit the module layer now that you have seen it in context.
• How neural networks learn — the same loop, told from the theory side.