Skip to content

← JAX Vectorization & Control Flow step 9 of 25

Medium Primitives

Scan with Per-Step Outputs

Jump to the task

Save one result from each step

lax.scan returns (final_carry, stacked_outputs). The carry holds state needed by the next step, such as a model's weight. The second return value collects results you want to keep, such as one loss per update. Unlike a Python loop unrolled during JIT tracing, scan keeps the iteration as a loop and stacks those results for you.

JAX stacks each array emitted in the body's second return value along a new leading dimension. None collects no output. If the body emits a scalar, the output is shape (N,). If it emits a vector of length d, the output is shape (N, d). The shape must be identical across all steps. Each output leaf must also keep its dtype and the output's structure must stay the same.

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
    value = loss(w, x, y)
    g = grad(w, x, y)
    new_w = w - 0.1*g
    return new_w, value

xs = jnp.ones((2, 2))
ys = jnp.ones((2, 2))
result = lax.scan(
    step, 0.0, (xs, ys)
)
final_w, losses = result
# final_w ≈ 0.64
# losses ≈ [2.0, 0.72]

Each row of xs and ys is one microbatch. Here there are two microbatches of two examples each. The loss is computed with the incoming weight w, before its update:

Step w in Loss w out
0 0.00 2.00 0.40
1 0.40 0.72 0.64

The weight crosses the step boundary as the carry; the loss is saved separately as that step's output. Unlike a cumulative sum, the output is not the new carry. losses has shape (2,), while final_w remains a scalar. The task below packs them into [0.64, 2.0, 0.72].

Common pitfalls

  • Loss after update. Computing the loss with new_w instead of w gives a different measurement: the updated model evaluated on the same batch. That is valid elsewhere, but this exercise asks for the pre-update loss. It is not generally the next batch's loss.
  • Inconsistent output shape. If your body emits different shapes on different iterations, scan will error at trace time. Keep it static.
  • Forgetting jnp.atleast_1d on the scalar carry. final_w is a 0-D array; jnp.concatenate requires at least 1-D inputs, so wrap it with jnp.atleast_1d before concatenating.

Losses from different batches need not decrease, because they measure different examples. Even on one batch, too large a learning rate can increase the loss. Keep the recorded values in batch order rather than sorting them or treating them as a guaranteed progress curve.

You may use jax.value_and_grad(loss, argnums=0) to obtain the loss and its w derivative together. Both must use the incoming weight.

Problem

Implement scan_collect_losses(w, lr, x_batches, y_batches) using lax.scan. The batch arrays share shape (N, B), where N and B are positive. The weight w and learning rate lr are scalars. Use real floating-point inputs with matching dtypes so the carry's dtype stays unchanged.

For each paired row (x, y), compute sum((w*x - y)**2) and its derivative with respect to the current w. Record that loss, then update w = w - lr*gradient. Use the updated weight on the next row. Keep lr fixed and use a sum, not a mean, over each batch.

Return a one-dimensional array of shape (N+1,): [final_w, loss_step_0, ..., loss_step_{N-1}]. The leading entry is the final weight, not the final loss. Do not append an extra loss evaluated after the last update.

One way to pack the result is jnp.concatenate([jnp.atleast_1d(final_w), losses]). Output tests check the values and ordering; they cannot enforce scan over a numerically equivalent loop.