NNX Checkpoint / Remat

A forward function can accept an NNX model and call it directly. nnx.remat wraps that function with gradient checkpointing: during differentiation, some values needed for derivatives can be recomputed instead of retained from the forward pass.

Build the model first

class MLP(nnx.Module):
    def __init__(self, num_layers, features, *, rngs):
        self.layers = nnx.List([
            nnx.Linear(features, features, rngs=rngs)
            for _ in range(num_layers)
        ])

    def __call__(self, x):
        for layer in self.layers:
            x = jax.nn.relu(layer(x))
        return x

Use one nnx.Rngs(seed) object for the constructors, in order. Each Linear consumes fresh keys from that stream. The layers have separate parameters; their default zero biases can still have equal values. Creating a new RNG object with the same seed for every layer would repeat initialization instead of advancing the stream.

nnx.List is a state-aware container for these submodules. Ordinary lists are valid JAX pytrees, so the reason for this container is NNX attribute and state handling, not a general ban on lists in pytrees. Explicitly marking a list with nnx.data is another supported approach in current NNX; use nnx.List for this exercise.

Wrap the forward function

@nnx.remat
def fwd(model, x):
    return model(x)

model = MLP(N, F, rngs=nnx.Rngs(seed))
output = fwd(model, x)

The NNX transformation handles module state so the wrapped function can accept the model object. Initialization happens before this call. The model here has no dropout, running statistics or optimizer updates.

Checkpointing preserves the mathematical output and derivatives. Compare floating-point results with a tolerance rather than expecting byte-for-byte equality across compilation settings. A forward-only call does not measure backward-pass memory savings.

Values saved for differentiation are called residuals. Checkpointing can reduce their storage by recomputing internal values from retained inputs. It does not remove all activation storage or guarantee constant memory in network depth. Cost depends on the boundary, model, compiler and hardware; there is no fixed extra-compute percentage. Profile a differentiated workload when choosing boundaries or nested checkpoint regions.