Skip to content

← Flax NNX step 53 of 100

Medium Primitives

NNX Eval Step

Jump to the task

Why this matters

Training changes parameters using a loss on the training data. Evaluation measures predictions on another dataset without using it for parameter updates. Here you will train a Linear for three SGD steps, then report two error measures on the evaluation data.

Use the trained parameter values for evaluation. Reusing the same model is convenient, but a separate model loaded with the same trained parameters can also be valid. Constructing a fresh random model would evaluate different weights.

Two ways to measure error

For prediction errors pred - target, mean squared error (MSE) averages their squares, while mean absolute error (MAE) averages their magnitudes. Both means here cover every example and output coordinate. MSE has squared target units; MAE has the target's units.

import jax.numpy as jnp

pred = jnp.array([[2.0, -1.0], [0.0, 4.0]], dtype=jnp.float32)
target = jnp.array([[1.0, 1.0], [0.0, 1.0]], dtype=jnp.float32)
error = pred - target  # [[1, -2], [0, 3]]
mse = jnp.mean(error ** 2)
mae = jnp.mean(jnp.abs(error))
assert jnp.allclose(mse, 3.5)  # (1 + 4 + 0 + 9) / 4
assert jnp.allclose(mae, 1.5)  # (1 + 2 + 0 + 3) / 4

Check shapes before computing these metrics. In particular, targets shaped (N,) can broadcast unexpectedly against predictions shaped (N, 1). This exercise requires matching two-dimensional prediction and target shapes.

Evaluation behavior depends on the layers

This model contains only a Linear, whose forward pass does not update its parameters or use a training-mode flag. Evaluation therefore needs only a forward pass and metric calculations, with no optimizer update.

That does not generalize to every model. Dropout and BatchNorm need suitable evaluation settings; simply omitting the optimizer call does not disable stochastic behavior or running-statistic updates. Those layers are outside this exercise. Evaluation settings also do not, by themselves, prohibit differentiation: gradients are unnecessary for these metrics, but may be useful in other analysis tasks.

Problem

Implement eval_metrics(seed, x_train, y_train, x_eval, y_eval, lr):

  1. Construct nnx.Linear(x_train.shape[-1], y_train.shape[-1], rngs=nnx.Rngs(int(seed))) with default initializers and bias.
  2. Create one nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param). Use plain SGD, without momentum or weight decay.
  3. Perform exactly three updates on the entire training batch. For each, compute gradients of jnp.mean((model(x_train) - y_train) ** 2) at the current parameter values, then call optimizer.update(model, grads).
  4. After the third update, compute pred_eval = model(x_eval). Do not use evaluation targets in a training update.
  5. Compute MSE and MAE against y_eval and return jnp.stack([mse, mae]).

The optimizer persists across the three steps. Recompute gradients each time; do not reuse the gradients from the first step. Keep the learning rate and metrics as JAX values instead of converting them to Python floats. The seed must be static if JIT-compiling the entire entrypoint.

The returned values describe evaluation after training, not the training loss. Three steps do not guarantee good generalization or a lower evaluation error. The objective here is to implement the specified sequence accurately.

Inputs: x_train has shape (N_train, D_in) and y_train has shape (N_train, D_out). x_eval has shape (N_eval, D_in) and y_eval has shape (N_eval, D_out). All dimensions are positive; training and evaluation batch sizes may differ. Arrays are finite float32 values, small enough for finite intermediate calculations. seed is a nonnegative integer-valued scalar supplied as a float; cast it to int. lr is finite and nonnegative.

Output: JAX array of shape (2,), containing [mse_eval, mae_eval]. It is not necessarily stored on the host just because it contains two scalars.