Skip to content

← Flax NNX step 54 of 100

Medium Primitives

NNX MultiMetric

Jump to the task

Why this matters

To report metrics across several batches, you need to accumulate values without accidentally changing their weighting. nnx.MultiMetric groups named metric objects and forwards updates to them. Each child defines its own aggregation; the Average children used here track a sum and a count.

Match metric arguments to update keywords

import jax.numpy as jnp
from flax import nnx

metrics = nnx.MultiMetric(
    mse=nnx.metrics.Average("mse"),
    mae=nnx.metrics.Average("mae"),
)
metrics.update(mse=jnp.array(2.0), mae=jnp.array(1.0))
metrics.update(mse=jnp.array(4.0), mae=jnp.array(2.0))
result = metrics.compute()
assert jnp.allclose(result["mse"], 3.0)
assert jnp.allclose(result["mae"], 1.5)

metrics.reset()
metrics.update(mse=jnp.array(8.0), mae=jnp.array(3.0))
assert jnp.allclose(metrics.compute()["mse"], 8.0)

Average("mse") reads the mse keyword passed to update. The default argument name is values; naming a MultiMetric child mse does not itself change that child's argument name. compute() returns a dictionary of JAX scalar results. It does not reset the accumulated state; use reset() when starting a fresh measurement period.

Average batches or individual errors?

Passing one scalar mean per batch gives every batch equal weight. That equals the mean over all error entries when batches have equal sizes and matching output widths, as they do here.

For unequal batches, the difference can be large. A one-example batch with mean error 10 and a nine-example batch with mean error 0 have an unweighted batch average of 5, but an example-weighted average of 1. Do not assume the discrepancy is small. One alternative for Average is to pass the individual error values so its count reflects their number, rather than pre-averaging batches. This exercise specifically asks for scalar batch means.

Training data and reported metrics

This exercise trains and measures on the same supplied data. The results describe training-set error after two updates, not held-out performance. Splitting the data into measurement batches does not turn it into a separate evaluation dataset or establish generalization.

Keep the model unchanged while aggregating metrics. The Linear used here has no Dropout or BatchNorm state, so its forward calls do not require a train/eval flag. The metrics themselves hold changing state as they collect observations. With NNX transforms, pass stateful objects explicitly rather than assuming a captured Python object is an ordinary array leaf.

Problem

Implement multimetric_aggregate(seed, x_batches_flat, y_batches_flat, batch_size, lr):

  1. Build a default nnx.Linear from the input and target widths using nnx.Rngs(int(seed)). Create one plain-SGD optimizer with nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param).
  2. Perform exactly two updates on the full supplied dataset. Recompute gradients each time for jnp.mean((model(x) - y) ** 2).
  3. Construct the MultiMetric above, with an Average for each named metric.
  4. Slice both arrays into consecutive chunks of int(batch_size) rows. For each, compute scalar MSE and MAE over all rows and output coordinates, then call metrics.update(mse=mse, mae=mae).
  5. After processing all chunks, call result = metrics.compute() once and return jnp.stack([result["mse"], result["mae"]]).

Do not reset between chunks or update the model during aggregation. Keep learning rate and results as JAX values. If JIT-compiling the whole entrypoint, the seed and batch size must be static because they control initialization and Python slicing. The output tests cannot prove you used MultiMetric rather than a direct full-dataset mean; follow the requested aggregation workflow.

Inputs: x_batches_flat has shape (N, D_in) and y_batches_flat has shape (N, D_out), with positive dimensions and finite float32 values. batch_size is a positive integer-valued scalar that divides N exactly; no remainder is dropped. seed is nonnegative and integer-valued. Both scalars arrive as floats and are cast to int. lr is finite and nonnegative; values are small enough for finite intermediate computations.

Output: shape (2,), containing the mean of the batch MSEs and the mean of the batch MAEs, measured after the second training update.