Skip to content

← Flax step 61 of 100

Medium Primitives

eval_step - Forward + Metrics

Jump to the task

An evaluation step computes predictions and metrics from the current model parameters. It does not update those parameters or advance the training step. This exercise first trains a small linear model, then evaluates its final state.

Training and evaluation have different jobs

The training step computes mean squared error, differentiates it and applies an SGD update. The evaluation step only computes predictions, mean squared error (MSE) and mean absolute error (MAE):

preds = state.apply_fn({"params": state.params}, x).reshape(-1)
errors = preds - y
mse = jnp.mean(errors ** 2)
mae = jnp.mean(jnp.abs(errors))

Both metrics average over examples. MSE squares each error before averaging; MAE takes its absolute value. Neither is a sum, and MSE has no one-half factor. For errors [1, -3], MSE is 5 and MAE is 2.

The Dense layer returns shape (N, 1), while targets have shape (N,). Reshape predictions to (N,) before subtraction. Otherwise broadcasting can create an (N, N) matrix of errors between unrelated examples.

Keep evaluation read-only

TrainState holds parameters, optimizer state and a step counter. apply_gradients returns a new state after a training update. An evaluation helper should return metrics without calling it. Repeated evaluation of this Dense-only model with the same state and inputs gives the same results.

Models with dropout or BatchNorm need additional evaluation settings, such as disabling dropout and using stored BatchNorm statistics. There is no universal Linen training=False switch: the module's own call signature controls those options. This model has neither layer, so no extra state or random key is needed for evaluation.

The exercise evaluates on its training batch to keep the calculation small. Those metrics describe fit to that batch, not performance on unseen data.

Problem

Implement eval_step_metrics(seed, x, y):

  1. Define a Linen TinyDense module whose compact __call__ applies one unnamed nn.Dense(1) with default settings and bias. Keep this wrapper: its Dense_0 parameter path affects the seeded initialization.
  2. Initialize with jax.random.PRNGKey(int(seed)) and the full x, without splitting the key. Create a TrainState from the params collection, model.apply and optax.sgd(0.05) with its default settings.
  3. Take exactly three full-batch SGD steps on mean squared error using the same (x, y) each time. Update both the kernel and bias.
  4. Compute MSE and MAE from the parameters after the third update, using a read-only evaluation helper. Return them in that order as a JAX array of shape (2,).

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. Leave x and y unchanged. Do not return the loss computed before the third update; recompute predictions from the final parameters. The wrapper is not required to accept traced scalar arguments under jit.