NNX Remat Policies

A checkpoint policy lets selected intermediate results remain eligible for storage during differentiation. Saving an expensive-to-recompute result may avoid work; recomputing a cheap result may reduce storage. The tradeoff depends on result sizes and the surrounding computation.

What the policy controls

Pass the callable jax.checkpoint_policies.dots_with_no_batch_dims_saveable to nnx.remat. For dot_general operations, this policy permits saving the result when the operation has no paired batch dimensions. It does not force every eligible result to be stored.

A vector times a shared matrix, as used here, has no paired batch axes. So does (B, D) @ (D, F): the leading axis counts examples, but is not paired with an axis of the weight matrix. In contrast, independent products (B, N, D) @ (B, D, F) pair the two B axes. This rule is about operation structure, not whether a result is small.

Without an explicit policy, checkpointing does not retain internal residuals, and recomputes required values from saved inputs. The named policy can retain more results to reduce that work. Neither choice is universally faster or lower in peak memory.

Configure the decorator

@nnx.remat(
    policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable
)
def fwd(model, x):
    return model(x)

In this API, nnx.remat(policy=...) returns a decorator. The bare @nnx.remat form is also valid, but uses the default policy. The transformation understands NNX model objects and their state.

For this exercise, the model contains two layers:

class MLP(nnx.Module):
    def __init__(self, features, *, rngs):
        self.l1 = nnx.Linear(features, features, rngs=rngs)
        self.l2 = nnx.Linear(features, features, rngs=rngs)

    def __call__(self, x):
        x = jax.nn.relu(self.l1(x))
        return jax.nn.relu(self.l2(x))

Construct both layers with the same RNG stream, in that order. They receive fresh initialization keys and have separate parameters. Default biases are zero, so separate parameters need not have different values.

What the output can establish

The policy preserves the mathematical forward function and derivatives. Use tolerances for floating-point comparisons. Equal forward values do not establish that a policy was applied or measure its storage decisions. To investigate those decisions, differentiate a scalar loss and inspect residuals or profile the resulting computation. This model does not include dropout, running-statistic changes or a training update.