Skip to content

← Flax NNX step 52 of 100

Medium Primitives

NNX Train Step

Jump to the task

Why this matters

A reusable training step computes predictions, a scalar loss and its gradients, then updates parameters. In this exercise, the step receives an NNX model and optimizer explicitly. Updating them lets the next call continue from the new parameter values.

Construct the model and optimizer once, outside the loop. Reinitializing the model would discard learning; rebuilding an optimizer can reset state such as momentum. The plain SGD used here has no momentum, but the same reuse pattern works when the optimizer does carry it.

A complete example

import jax.numpy as jnp
import optax
from flax import nnx

model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.sgd(0.01), wrt=nnx.Param)
x = jnp.array([[1.0, 2.0], [-1.0, 0.5]], dtype=jnp.float32)
y = jnp.array([[1.0], [-0.5]], dtype=jnp.float32)

def loss_fn(model, x, y):
    return jnp.mean((model(x) - y) ** 2)

def train_step(model, optimizer, x, y):
    loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)
    optimizer.update(model, grads)
    return loss

losses = []
for _ in range(3):
    losses.append(train_step(model, optimizer, x, y))
print(jnp.stack(losses))

Each returned loss was computed before that step's update. Returning the scalar does not undo the update; the model retains its changed parameters. The function need not return the model or optimizer for this calling pattern.

Which loss is the last one?

Number parameter states by how many updates have been applied. If the initial state is p0, three calls record losses at p0, p1 and p2, while leaving the model at p3. The exercise returns the first and last recorded losses. It does not evaluate an additional loss after the final update.

For num_steps=1, both returned entries therefore contain the initial loss, even though the function must still perform one update. With five steps, the final recorded loss is measured after four updates, before the fifth.

Reusing one batch can reduce its loss, but monotonic decrease is not guaranteed for an arbitrary learning rate. Zero learning rate gives no parameter change, a zero gradient can leave loss unchanged, and a large rate can increase it. Follow the update rule rather than forcing the recorded losses to decrease.

Problem

Implement train_step_loss(seed, x, y, lr, num_steps):

  1. Construct nnx.Linear(x.shape[-1], y.shape[-1], rngs=nnx.Rngs(int(seed))) with default initializers and bias.
  2. Construct one nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param). Use plain SGD without momentum or weight decay.
  3. Define mean squared error as jnp.mean((model(x) - y) ** 2), averaging over every example and output coordinate.
  4. Define a step that obtains loss, grads from nnx.value_and_grad, calls optimizer.update(model, grads) once and returns that loss.
  5. Call the step int(num_steps) times on the same batch, recomputing gradients at the current parameters each time. Record each returned loss.
  6. Return jnp.stack([losses[0], losses[-1]]).

Keep learning rate and losses as JAX values, without Python float conversion. The seed and step count must be static if JIT-compiling the whole entrypoint. Numerical output tests do not verify the last update, since no returned loss depends on it. Still perform every requested step.

Inputs: x has shape (N, D_in) and y has shape (N, D_out), with matching positive batch size and positive widths. Both are finite float32 arrays. seed is a nonnegative integer-valued scalar and num_steps is a positive integer-valued scalar, supplied as floats and cast to int. At least one step is required to record a loss. lr is finite and nonnegative; values are small enough to keep intermediate computations finite.

Output: shape (2,), containing the loss before the first update and the loss before the final update.