Skip to content

← Flax step 59 of 100

Medium Primitives

TrainState - One Step

Jump to the task

A training step needs the current parameters and the optimizer's state. flax.training.train_state.TrainState keeps these together with the model's apply function and a step counter. This exercise uses it to take one SGD step on a small Flax Linen model.

Create the state

Use keyword arguments:

state = train_state.TrainState.create(
    apply_fn=model.apply,
    params=variables["params"],
    tx=optax.sgd(lr),
)

params holds the parameter collection, while apply_fn expects a full variables mapping. For this Dense-only model, call it with {"params": state.params}. The state also stores the optimizer transformation as tx, initializes opt_state by calling tx.init(params), and starts step at zero.

Default SGD has no momentum buffers. Other optimizers may keep additional arrays in opt_state; TrainState passes that state from one update to the next.

One update returns a new state

Compute gradients with respect to the parameter tree, then call:

new_state = state.apply_gradients(grads=grads)

This uses the optimizer to turn gradients into updates, applies them to the parameters and returns a state with step increased by one. It does not compute the gradients for you. The original state is unchanged by this method, so use the returned state for later work.

TrainState's fields cannot be reassigned directly, but that does not recursively freeze a regular Python dictionary stored in params. Treat the parameter tree as immutable and use functional updates instead of editing nested dictionaries.

Measure the updated model

For predictions preds and targets y, use mean((preds - y)**2) with no one-half factor. A Dense layer with one output feature produces shape (N, 1); reshape it to (N,) before subtracting the targets. Otherwise broadcasting would compare predictions against unrelated examples.

Recompute predictions after applying the gradients. The loss used to obtain the gradients describes the old parameters, not the updated ones. A large learning rate can increase the loss, so do not assume the final value is lower.

Problem

Implement train_state_step_count(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. The wrapper's Dense_0 parameter path matters for seeded initialization.
  2. Initialize it with jax.random.PRNGKey(int(seed)) and x, without splitting the key. Create a TrainState with the params collection, model.apply and optax.sgd(float(lr)) using default settings.
  3. Differentiate the full-batch mean squared error with respect to both the kernel and bias. Call apply_gradients(grads=grads) exactly once.
  4. Recompute MSE using the returned state's parameters. Return a floating-point JAX array of shape (2,) containing [step_after, loss_after].

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 the input arrays unchanged.

The first output is 1.0, even when lr=0: the update call advances the counter although the parameters do not change. In that case the second output is the initialized model's MSE.