The Complete Training Loop

Let's put everything together into a complete, production-style training loop with Flax NNX and Optax. This is the pattern you'll use (or adapt) for real projects.

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

# --- Model Definition ---
class Classifier(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
        self.linear1 = nnx.Linear(64 * 5 * 5, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, 10, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)

    def __call__(self, x):
        x = nnx.relu(self.conv1(x))
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nnx.relu(self.conv2(x))
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape(x.shape[0], -1)
        x = nnx.relu(self.linear1(x))
        x = self.dropout(x)
        return self.linear2(x)

# --- Setup ---
model = Classifier(rngs=nnx.Rngs(0))

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=1e-3,
    warmup_steps=500,
    decay_steps=10000,
)

tx = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule),
)
optimizer = nnx.Optimizer(model, tx)

# --- Training Step ---
@nnx.jit
def train_step(model, optimizer, x, y):
    def loss_fn(model):
        logits = model(x)
        loss = optax.softmax_cross_entropy_with_integer_labels(logits, y)
        return jnp.mean(loss)

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(grads)
    return loss

# --- Evaluation Step ---
@nnx.jit
def eval_step(model, x, y):
    logits = model(x)
    predictions = jnp.argmax(logits, axis=-1)
    accuracy = jnp.mean(predictions == y)
    return accuracy

# --- Training Loop ---
num_epochs = 10
for epoch in range(num_epochs):
    # Training
    train_loss = 0.0
    num_batches = 0
    for x_batch, y_batch in train_loader:
        loss = train_step(model, optimizer, x_batch, y_batch)
        train_loss += loss
        num_batches += 1

    avg_loss = train_loss / num_batches

    # Evaluation
    total_acc = 0.0
    eval_batches = 0
    for x_val, y_val in val_loader:
        acc = eval_step(model, x_val, y_val)
        total_acc += acc
        eval_batches += 1

    avg_acc = total_acc / eval_batches
    print(f"Epoch {epoch}: loss={avg_loss:.4f}, val_acc={avg_acc:.4f}")