Skip to content

← JAX Autodiff step 20 of 25

Hard Primitives

Microbatched Gradient Accumulation via scan

Jump to the task

Add gradients before updating the weight

Gradient accumulation splits a batch into smaller groups called microbatches. Evaluate each group at the same parameter values, add the gradients, then use the total for one optimizer update. Updating the weight between groups would be a different training procedure.

For a loss that sums independent example losses, this total agrees with the full-batch gradient up to floating-point rounding. The example below uses squared errors and a fixed weight of one:

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

xs = jnp.array([[1.0, 2.0], [3.0, 4.0]])
ys = jnp.zeros_like(xs)
w = jnp.array(1.0)

def loss(w, x, y):
    return jnp.sum((w*x - y)**2)

grad_loss = jax.grad(loss)
def step(acc, batch):
    x, y = batch
    return acc + grad_loss(w, x, y), None

total, _ = lax.scan(step, jnp.zeros_like(w), (xs, ys))
# total = 60.0

The first group contributes 2*(1**2 + 2**2) = 10; the second contributes 2*(3**2 + 4**2) = 50. Their sum is 60. For a mean loss over all four examples, divide by four to get 15. Summing microbatch means instead would give 30, not the full-batch mean gradient.

Carry the running total through scan

Reshape the data from (N,) to (n_micro, micro_size). lax.scan reads one row from each data array per step. The body receives the current accumulator and the paired rows, then returns the new accumulator and a per-step output. Use None for that output when only the final total is needed.

The carry must keep the same shape and dtype at every step. A scalar weight has a scalar gradient, so initialize with jnp.zeros_like(w). A length-one vector is a different shape from a scalar.

The diagram adds the already-computed gradients 10 and 50. It shows running totals for illustration; the Python code returns only the final carry. The diagram does not compute gradients or measure memory.

Loading visualization…

What splitting changes

Smaller groups can reduce intermediate activation storage, but the parameters, gradient accumulator and full input arrays still occupy memory. Do not infer total memory use from microbatch size alone. Also, splitting must preserve the objective: training-time batch normalization, for example, can use different statistics after a split.

The reshape needs a concrete microbatch count. When JIT-compiling this exercise, use jax.jit(microbatched_grads, static_argnames="n_microbatches"). Calling int on a traced value does not make that value static.

Problem

Implement microbatched_grads(w, b, x_full, y_full, n_microbatches). Split both data arrays into equal contiguous chunks. Use lax.scan to accumulate the gradient with respect to w of each chunk's loss sum((w*x - y)**2), keeping w fixed.

w and b are real floating-point scalars. x_full and y_full are real floating-point JAX arrays of shape (N,), with matching dtypes. The positive whole-number count arrives as a float, such as 2.0; convert it to an integer. N is positive and divisible by that count. There is no partial final chunk in this exercise.

Return the summed scalar gradient. Do not average it or update w. The argument b is deliberately unused, even if it is nonzero. For w=1, zero targets, inputs [1, 2, 3, 4] and two microbatches, return 60 as in the example above.