Skip to content

← JAX Autodiff step 17 of 25

Medium Primitives

Checkpointed Deep Stack via scan

Jump to the task

A stack of layers repeatedly updates an activation array. jax.lax.scan expresses that repetition with a body function and a carried value. Here the carry is the current activation, and each step reads one weight matrix from the leading axis of weights.

Define one layer

import jax
import jax.numpy as jnp
from jax import lax

def layer(y, w):
    return jax.nn.relu(y @ w), None

The first returned item becomes the next step's carry. The second would normally be collected as a per-step output. Returning None avoids constructing that output stack; it does not prevent autodiff from keeping values needed for derivatives.

Every step must preserve the carry's shape and dtype. With float32 activations of shape (N_batch, d) and float32 weights of shape (d, d), this body does so. Each layer has its own matrix; equal initial values do not make the layers share one parameter.

Differentiate the final loss

checkpointed_layer = jax.checkpoint(layer)

def loss(weights, x):
    final, _ = lax.scan(checkpointed_layer, x, weights)
    return jnp.sum(final ** 2)

weights = jnp.stack([jnp.eye(2)] * 3)
x = jnp.ones((1, 2))
gradients = jax.grad(loss, argnums=0)(weights, x)
# Shape: (3, 2, 2). Every entry is 2.0.

The identity matrices leave [1, 1] unchanged. For each layer, its input is [1, 1] and the derivative flowing back is [2, 2]. Their outer product is the matrix filled with 2. The returned array contains a separate gradient matrix for every layer.

What checkpointing changes

Reverse-mode differentiation uses saved values called residuals. Checkpointing the body lets JAX recompute internal residuals when needed instead of retaining them across the forward and backward passes. It preserves the mathematical loss and gradient; compare floating-point results with a tolerance.

This does not guarantee constant memory as the number of layers grows. Layer inputs, parameters and other residuals can still grow with depth. Inspect residuals and measure an actual differentiated workload to assess memory and runtime. Returning None from the scan body is not itself gradient checkpointing.

Transform the body function with jax.checkpoint(layer), not the result tuple of lax.scan. Wrapping a function containing a whole scan is also valid, but chooses a different checkpoint boundary.

Problem

Implement grad_checkpointed_stack(weights, x) for eight layers:

  1. Define the body above and wrap it with jax.checkpoint.
  2. Scan over weights from index 0 through 7, starting with carry x.
  3. Sum the squares of the final activation across all batch and feature entries. Apply ReLU at every layer, including the last.
  4. Return the gradient of that scalar loss with respect to weights.

weights is a finite float32 JAX array of shape (8, d, d), and x is a finite float32 JAX array of shape (N_batch, d), with both dimensions positive. There are no biases. Leave both inputs unchanged.

Return a JAX array of shape (8, d, d), not the scalar loss, the input gradient, or updated weights. Use a sum rather than a mean. At a zero preactivation, follow jax.nn.relu's zero-derivative convention.

Numerical tests check the weight gradients. They cannot prove that a submission used scan or checkpointing, and do not measure memory savings.