Skip to content

← JAX Autodiff step 16 of 25

Hard Primitives

Checkpoint with Save Policy

Jump to the task

A checkpoint policy controls which intermediate results may be retained for differentiation. Keeping a costly matrix-product result can avoid recomputing it later. Recomputing a cheap operation can reduce storage. The right choice depends on result sizes, computation cost and the rest of the model.

Saving is permission, not a command

jax.checkpoint(f) normally prevents internal derivative residuals from being saved. A residual is a value needed to calculate a derivative. Required values are recomputed from saved inputs instead. Operations whose results are unnecessary for the derivative need not run again.

A standard saving policy receives a primitive operation, abstract argument information and operation parameters. It returns a boolean indicating whether the result is eligible to be saved as a residual. It does not inspect the actual training data or force every eligible result to be kept.

  • nothing_saveable allows no internal residuals to be saved, matching checkpointing without an explicit policy.
  • everything_saveable allows all intermediate results to be saved if needed, following the ordinary residual-saving strategy.
  • dots_with_no_batch_dims_saveable allows results of dot_general operations with no paired batch dimensions to be saved.

Which batch dimensions?

In this exercise, (N, d) @ (d, d) multiplies each input row by the same weight matrix. Its dot operation has a contraction dimension but no paired batch dimensions, even though N counts examples.

By contrast, (B, N, d) @ (B, d, F) uses a separate weight matrix for each batch element. That dot pairs the B axes and is not selected by this rule. The rule does not classify outputs as small or large.

Apply the policy

Define a pure function containing the full loss computation, then wrap it:

def loss(w1, w2, x):
    hidden = jax.nn.relu(x @ w1)
    output = hidden @ w2
    return jnp.sum(output ** 2)

policy = jax.checkpoint_policies.dots_with_no_batch_dims_saveable
checkpointed_loss = jax.checkpoint(loss, policy=policy)
result = checkpointed_loss(w1, w2, x)

Pass the callable policy, not its name as a string. The wrapped function accepts the same arguments as loss. Its mathematical result and derivatives are unchanged; compare floating-point results with a tolerance. To investigate storage, differentiate the loss and inspect saved residuals or profile a training step. A forward-only result does not measure memory.

This policy is one option, not a universal production default. It allows more residuals to be retained than the default checkpoint policy, which may reduce recomputation but can increase storage. Neither lower peak memory nor a particular speedup follows from choosing it.

Problem

Implement checkpoint_dot_policy(w1, w2, x) using the loss above. Wrap the entire loss function in jax.checkpoint with jax.checkpoint_policies.dots_with_no_batch_dims_saveable.

w1 and w2 are finite float32 JAX matrices of shape (d, d). x is a finite float32 JAX matrix of shape (N, d), with N, d >= 1. There are no biases. Apply ReLU only after the first multiplication. Leave all inputs unchanged.

Return a scalar JAX array: the sum of squared final outputs across all examples and features. Do not average, apply a final ReLU, or return a gradient. With identity weights and x = [[1, 1]], the loss is 2.

Output tests check the numerical loss, but cannot prove which checkpoint policy was used. Follow the wrapping requirement even though a different policy, or no checkpoint wrapper, can produce the same loss.