Skip to content

← Flax NNX step 51 of 100

Medium Primitives

NNX Optimizer Basics

Jump to the task

Why this matters

A training step computes a loss, differentiates it with respect to the parameters, then applies an update. nnx.Optimizer connects an Optax update rule to selected NNX variables. You pass the model to optimizer.update; the method updates the selected parameters and the optimizer's own state.

In the API used here, keep the model and optimizer as separate objects. Do not rely on the optimizer storing your model. Its wrt=nnx.Param filter selects parameters for optimization. Other state, such as running statistics, may be updated by a model's forward pass instead.

A complete training step

import jax.numpy as jnp
import optax
from flax import nnx

model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.sgd(0.01), wrt=nnx.Param)
x = jnp.array([[1.0, 2.0], [-1.0, 0.5]], dtype=jnp.float32)
y = jnp.array([[1.0], [-0.5]], dtype=jnp.float32)

def loss_fn(model, x, y):
    return jnp.mean((model(x) - y) ** 2)

loss_before, grads = nnx.value_and_grad(loss_fn)(model, x, y)
optimizer.update(model, grads)
loss_after = loss_fn(model, x, y)
print(jnp.stack([loss_before, loss_after]))

value_and_grad returns the loss and gradients from the same forward evaluation. For this model argument, it differentiates Param variables by default. The returned gradient state contains entries for the kernel and bias. The optimizer then applies plain SGD to both: subtract the learning rate times each gradient.

Call optimizer.update for its effect on the model; do not assign its return value to the model. Return behavior differs across Flax versions. JAX arrays remain immutable, while NNX updates the values held by its Variable wrappers. In a multi-step loop, construct the model and optimizer once and reuse them, so optimizer state persists.

Before and after are different measurements

Report the loss returned by value_and_grad as the pre-step loss. Evaluate the loss again after the update to obtain the post-step loss.

A step need not strictly reduce loss. A zero learning rate leaves it unchanged, as does a zero gradient; an excessively large rate can increase it. The provided examples use rates that reduce their losses, but the exercise is about applying the specified update, not enforcing a decrease.

Problem

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

  1. Create nnx.Linear(x.shape[-1], y.shape[-1], rngs=nnx.Rngs(int(seed))) with default initializers and bias.
  2. Create nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param). Use no momentum or weight decay.
  3. Define mean squared error as jnp.mean((model(x) - y) ** 2), averaging over all batch and output entries.
  4. Obtain loss_before, grads with nnx.value_and_grad, then call optimizer.update(model, grads) once.
  5. Evaluate loss_after on the updated model using the same inputs.
  6. Return jnp.stack([loss_before, loss_after]).

Keep learning rate and losses as JAX values. Converting traced values to Python floats prevents this function from working under JIT. The seed sets up initialization and must be static if JIT-compiling the whole entrypoint.

Inputs: x has shape (N, D_in) and y has shape (N, D_out), with matching positive batch size and positive widths. Both are finite float32 arrays. seed is a nonnegative integer-valued scalar supplied as a float; cast it to int. lr is a finite nonnegative scalar. Values are small enough for finite loss, gradient and update calculations.

Output: array of shape (2,), containing [loss_before, loss_after].