TopolearnThis pattern is the backbone of every JAX training script, including those built on Flax and Optax.
In JAX, training state is data you pass around, not objects you mutate. A training step is a pure function:
new_params, new_opt_state, loss = train_step(params, opt_state, batch)Inside it you compute jax.value_and_grad(loss_fn)(params, batch), apply an update with jax.tree.map, and return everything new. Decorate it with @jax.jit so the whole step, forward and backward pass plus update, compiles to one executable.
Patterns to follow:
- Keep parameters in a pytree (dict or registered dataclass); initialise them from split random keys.
- Thread the RNG key through the step if you use dropout, splitting each step.
- Iterate in Python over batches and call the jitted step; the Python loop is cheap compared to the compiled step.
- Keep batch shapes fixed (drop or pad the last batch) to avoid recompiling.
Real projects rarely hand-roll optimizers. Optax provides optax.adam(lr) with an init(params) and update(grads, state, params) API that fits this exact shape, and Flax (its nnx API) or Equinox provide layers. Learning the bare pattern first makes those libraries transparent: they are pytrees plus pure functions.
Watch for host-device syncs in the loop: calling float(loss) every step forces the program to wait for the device. Log every N steps instead.
Example
import jax
import jax.numpy as jnp
def init(key, n_in):
return {"w": jax.random.normal(key, (n_in,)) * 0.1, "b": jnp.zeros(())}
def loss_fn(params, x, y):
pred = x @ params["w"] + params["b"]
return jnp.mean((pred - y) ** 2)
@jax.jit
def train_step(params, x, y, lr=0.1):
loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
params = jax.tree.map(lambda p, g: p - lr * g, params, grads)
return params, loss
key, k_data, k_init = jax.random.split(jax.random.key(0), 3)
x = jax.random.normal(k_data, (256, 3))
y = x @ jnp.array([2.0, -1.0, 0.5]) + 3.0
params = init(k_init, 3)
for step in range(200):
params, loss = train_step(params, x, y)
print(jnp.round(params["w"], 2), jnp.round(params["b"], 2)) # [ 2. -1. 0.5] 3.0
print(float(loss) < 1e-4) # TrueRemember
- A train step is a pure function from (params, batch) to (new params, metrics)
- Jit the whole step; loop over batches in Python
- Optax and Flax follow the same pure, pytree-based pattern
Watch out for
- Converting the loss to a Python float every step, which blocks on the device and slows training.