Skip to content

← Flax step 60 of 100

Medium Primitives

train_step with value_and_grad

Jump to the task

A training step often needs both the loss for logging and its gradients for an optimizer update. jax.value_and_grad returns these together. This lesson uses it inside a helper that accepts a TrainState and a batch, then returns the updated state and the loss measured before the update.

Loss and gradient refer to the same parameters

loss, grads = jax.value_and_grad(loss_fn)(state.params, batch)
new_state = state.apply_gradients(grads=grads)
return new_state, loss

The returned loss is loss_fn(state.params, batch), evaluated at the input parameters. Updating the state does not change that already computed value. To measure the new parameters, you would need another forward pass. This exercise specifically returns the pre-update loss.

value_and_grad differentiates the first positional argument by default. Put params first and batch second. The loss must be a scalar array of shape (), and the resulting gradients have the same tree structure as the parameters. jax.grad returns only the gradients; value_and_grad also makes the loss available without a separate loss call. Actual performance depends on the computation and compilation, so this is not a guaranteed speedup.

Define the loss with explicit parameter inputs

Use mean squared error for the entire batch:

def loss_fn(params, batch):
    bx, by = batch
    preds = model.apply({"params": params}, bx).reshape(-1)
    return jnp.mean((preds - by) ** 2)

Capture the model's apply function, while passing parameters and batch data as arguments. This keeps the values being differentiated explicit.

The Dense layer returns (N, 1), but the targets have shape (N,). Reshape predictions before subtraction; otherwise broadcasting creates an (N, N) error matrix. The loss is the mean of squared errors, without a one-half factor. Both the kernel and bias must receive gradients.

Problem

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

  1. Define a Linen TinyDense module with a compact __call__ that applies one unnamed nn.Dense(1) using default settings and bias. Keep this wrapper so the parameter path is Dense_0, as required by the seeded outputs.
  2. Initialize it with jax.random.PRNGKey(int(seed)) and x, without splitting the key. Create a TrainState using model.apply, the initialized params collection and optax.sgd(float(lr)) with default settings.
  3. Define loss_fn(params, batch) as above. In train_step(state, batch), compute (loss, grads) with jax.value_and_grad, call state.apply_gradients(grads=grads) once and return (new_state, loss).
  4. Call that helper once on (x, y). Return its loss in a floating-point JAX array of shape (1,).

x is a finite float32 array of shape (N, D) and y is a finite float32 array of shape (N,), with N >= 1 and D >= 1. seed is a nonnegative integer-valued scalar that may arrive as a float; lr is finite and nonnegative. Leave both input arrays unchanged.

The updated state's step counter is one. Changing lr affects its parameters, but not the loss returned by this exercise, which was computed beforehand. Still perform the requested training update. The outer wrapper is not required to accept traced scalar arguments under jit.