nnx_optimizer

Chapter 8: nnx.Optimizer

In the previous chapter, we examined the internal GraphDef and GraphState representations used by the NNX Functional API (split/merge/state/update/graphdef). Now, we'll look at a practical helper class built upon these concepts that significantly simplifies the standard training loop structure in NNX: nnx.Optimizer.

Motivation: Simplifying Stateful Training Loops

A typical JAX training loop involves managing several pieces of state: the model's parameters, the optimizer's internal state (e.g., momentum vectors), and potentially a step counter. While you can manage these separately using the functional API or standard Python variables, it often leads to boilerplate code for passing state around and applying updates.

Flax Linen introduced flax.training.TrainState to bundle these components. nnx.Optimizer serves a similar purpose in the NNX ecosystem. It's a stateful container, inheriting from nnx.Object, designed to hold an nnx.Module instance, the corresponding Optax optimizer state (opt_state), and the Optax gradient transformation (tx). This colocation simplifies the training process by providing a single object that manages both the model's learnable parameters and the optimizer's evolving state. Its primary update() method streamlines the application of gradients.

Central Use Case: A Standard Training Step

Let's illustrate how nnx.Optimizer simplifies a basic training step involving gradient computation and parameter updates.

import jax
import jax.numpy as jnp
from flax import nnx
import optax # Standard JAX optimizer library

class SimpleModel(nnx.Module):
  def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(din, dout, rngs=rngs)

  def __call__(self, x: jax.Array) -> jax.Array:
    return self.linear(x)

# --- Setup ---
key = jax.random.key(0)
model_key, data_key = jax.random.split(key)
model = SimpleModel(din=10, dout=5, rngs=nnx.Rngs(model_key))

# Define Optax optimizer
tx = optax.adam(learning_rate=1e-3)

# 1. Instantiate nnx.Optimizer
#    It holds the model, the optimizer transformation (tx), 
#    and initializes the optimizer state internally.
optimizer = nnx.Optimizer(model, tx) 

x_batch = jax.random.normal(data_key, (32, 10))
y_batch = jnp.ones((32, 5)) # Dummy targets

# --- Loss Function ---
# Takes the *model* part of the optimizer as input
def loss_fn(model: SimpleModel, x: jax.Array, y: jax.Array):
  y_pred = model(x)
  loss = jnp.mean((y_pred - y) ** 2)
  return loss

# --- JITted Training Step ---
@nnx.jit # Use nnx.jit with nnx.Optimizer
def train_step(optimizer: nnx.Optimizer, x: jax.Array, y: jax.Array):
  # 2. Calculate gradients w.r.t. model's Params (default)
  #    Pass optimizer.model to the loss function
  grad_fn = nnx.value_and_grad(loss_fn)
  loss, grads = grad_fn(optimizer.model, x, y) 
  
  # 3. Apply gradients using optimizer.update()
  #    This updates optimizer.model parameters and optimizer.opt_state in-place.
  optimizer.update(grads)
  
  # optimizer object with updated state is returned by nnx.jit
  return loss, optimizer 

# --- Execute ---
print(f"Initial loss: {loss_fn(optimizer.model, x_batch, y_batch)}")

# Run one training step
# Note: nnx.jit returns a *new* optimizer instance with updated state
loss, updated_optimizer = train_step(optimizer, x_batch, y_batch) 

print(f"Loss after step: {loss}")
# Access the updated model via the returned optimizer
print(f"Loss after update (recomputed): {loss_fn(updated_optimizer.model, x_batch, y_batch)}")