Scan for Efficient Training and Checkpointing with Orbax

jax.lax.scan is JAX's way of writing efficient loops that compile to a single XLA operation. Instead of unrolling a Python loop (which creates one copy of the computation per iteration), scan reuses the same compiled body. This is especially useful for training loops that run inside JIT.

import jax
import jax.numpy as jnp

# Instead of a Python loop for multiple gradient steps:
# for _ in range(100):
#     params, loss = train_step(params, batch)

# Use scan for a compiled loop:
def scan_train(params, opt_state, batches):
    """Run multiple training steps as a single compiled scan."""
    def step_fn(carry, batch):
        params, opt_state = carry
        x, y = batch

        def loss_fn(p):
            logits = forward(p, x)
            return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(
                logits, y))

        loss, grads = jax.value_and_grad(loss_fn)(params)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return (params, opt_state), loss

    (final_params, final_opt_state), losses = jax.lax.scan(
        step_fn,
        init=(params, opt_state),
        xs=batches,  # stacked batch arrays: (num_steps, batch_size, ...)
    )
    return final_params, final_opt_state, losses

# This compiles to a SINGLE XLA while loop — much faster than
# running train_step in a Python for loop inside jit

Checkpointing with Orbax

Saving and restoring model parameters is essential for long training runs. Orbax is the official checkpointing library for JAX:

import orbax.checkpoint as ocp

# Save a checkpoint
checkpointer = ocp.StandardCheckpointer()

# Save parameters (works with any pytree)
save_path = '/tmp/my_model/step_1000'
checkpointer.save(save_path, params)

# Restore parameters
restored_params = checkpointer.restore(save_path)

# For Flax NNX models, save the state:
state = nnx.state(model)
checkpointer.save('/tmp/my_model/step_2000', state)

# Restore:
restored_state = checkpointer.restore('/tmp/my_model/step_2000')
nnx.update(model, restored_state)

# CheckpointManager for managing multiple checkpoints
manager = ocp.CheckpointManager(
    '/tmp/checkpoints',
    options=ocp.CheckpointManagerOptions(
        max_to_keep=3,           # keep only 3 most recent
        save_interval_steps=500, # save every 500 steps
    ),
)

# In training loop:
for step in range(10000):
    params, loss = train_step(params, batch)
    manager.save(step, args=ocp.args.StandardSave(params))

# Restore latest:
step = manager.latest_step()
params = manager.restore(step)