Skip to content

← JAX Vectorization & Control Flow step 8 of 25

Hard Primitives

Training Loop via lax.scan

Jump to the task

Pass the updated weight to the next batch

lax.scan is useful when a sequence of training updates can be expressed as a fixed-length loop over stacked microbatches. It lowers the loop to a compiled loop operation, avoiding the large unrolled computation a Python for loop can create inside jit. A Python loop around a jitted training step is also valid, especially when the host handles data loading or logging between updates.

The carry is the weight (or a parameter pytree), and xs contains microbatches stacked along axis 0. Each step consumes one microbatch, computes a gradient and passes the updated weight to the next step.

Worked mini-example

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

def loss(w, x, y):
    error = w*x - y
    squares = error**2
    return squares.sum()

grad = jax.grad(loss)

def step(w, batch):
    x, y = batch
    g = grad(w, x, y)
    new_w = w - 0.1*g
    return new_w, None

xs = jnp.ones((2, 2))
ys = jnp.ones((2, 2))
result = lax.scan(
    step, 0.0, (xs, ys)
)
final_w, outputs = result
# final_w ≈ 0.64
# outputs is None

Starting at w = 0, each microbatch contains two examples with input and target equal to one. The update is w - 0.1 * gradient:

Step Gradient New w
0 -4.00 0.40
1 -2.40 0.64

Passing (xs, ys) scans both arrays together: the body receives (xs[i], ys[i]) at step i. The scalar weight is the carry; returning None for the second result means no per-step output array is collected.

These are sequential updates, not gradient accumulation at a fixed weight. In the example, adding both gradients evaluated at w=0 would give w=0.8 instead of 0.64. The second gradient must use w=0.4. With different batches, changing their order can change the final weight.

The loss is a sum over each batch. Replacing it with a mean divides that batch's gradient by its size, changing the update at the same learning rate. Gradient descent also does not guarantee a minimum after N updates; a large learning rate can make the loss increase.

Keep the scan state consistent

  • Forgetting to unpack the batch. Inside the body, batch is a tuple (x, y). Unpack it with x, y = batch (or index its two arrays) before evaluating the loss; the pair itself is not one input array.
  • Learning rate versus carry. The body can close over lr from the enclosing function. It stays unchanged during this scan, so it need not be part of the carry. An lr argument to a jitted outer function can still be a dynamic value; closing over it in the scan body does not make it a static JIT argument.
  • Wrong per-step output. When you don't need per-step outputs, return None as the second element. lax.scan still needs the 2-tuple (new_carry, output).

The carry must keep the same structure, shape and dtype across steps. Here it is always one floating-point scalar. The batch arrays are read in order along their first axis; each row has B examples.

Problem

Implement scan_training_step(w, lr, x_batches, y_batches) with lax.scan. Starting from w, take one gradient-descent step per batch using the loss sum((w*x - y)**2):

gradient = 2 * sum((w*x - y) * x)
new_w = w - lr * gradient

The two batch arrays have the same shape (N, B), with positive N and B. Inputs are real floating-point values; use matching dtypes so each update preserves the scalar carry's dtype. The scalar lr is the learning rate and stays fixed throughout the scan. It need not be a static argument when you JIT the enclosing function.

Return only the scalar final weight after all N updates. Use each paired row (x_batches[i], y_batches[i]) exactly once, in order. Do not reuse the first row, sum gradients at the initial weight, or collect a history of weights as the return value. The output tests check these updates, but cannot enforce use of scan over an equivalent Python loop.