One trace, however many steps

Here is a linear model taking one gradient step per batch, scanned over six batches. The whole loop is one equation, and its params tell you how to read the four operands after it. num_consts=1 claims the first, num_carry=1 claims the second, and everything left is sliced along its leading axis, one row per step.

Line the operand list up against the body signature and the correspondence is exact. The operands are d a b c: the learning rate, the weights, the inputs, the targets. The body binders are g h i j: the learning rate whole, the weights as carry, and one slice each of inputs and targets, f32[2,3] and f32[2] cut from f32[6,2,3] and f32[6,2].

What decides that split is your closure, not a keyword. The learning rate is a traced argument closed over by the step function, so it became a constant operand. Pass 0.1 as a Python float instead and it stops being an operand at all: num_consts drops to 0 and the number is a literal inside the body.

The kernel path's jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → chapter teaches the nesting grammar this equation is written in, and the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → path's ingestion chapter walks one scan across the border into StableHLO. What this section adds is the reading of the params themselves, which is the part you need when the loop is your own training step and you want to know what got hoisted.

import jax
import jax.numpy as jnp

def train(w, xs, ys, lr):
    def step(w, batch):
        x, y = batch
        loss = lambda w: jnp.mean((x @ w - y) ** 2)
        return w - lr * jax.grad(loss)(w), loss(w)
    return jax.lax.scan(step, w, (xs, ys))

print(jax.make_jaxpr(train)(jnp.zeros(3), jnp.ones((6, 2, 3)), jnp.ones((6, 2)), jnp.float32(0.1)))

# { lambda ; a:f32[3] b:f32[6,2,3] c:f32[6,2] d:f32[]. let
#     e:f32[3] f:f32[6] = scan[
#       _split_transpose=False
#       jaxpr={ lambda ; g:f32[] h:f32[3] i:f32[2,3] j:f32[2]. let
#           k:f32[2] = dot_general[
#             dimension_numbers=(([1], [0]), ([], []))
#             preferred_element_type=float32
#           ] i h
#
#           ... 16 more equations: the loss, the gradient, the update ...
#
#           ba:f32[] = div z 2.0
#         in (v, ba) }
#       length=6
#       linear=(False, False, False, False)
#       num_carry=1
#       num_consts=1
#       reverse=False
#       unroll=1
#     ] d a b c
#   in (e, f) }

import jax
import jax.numpy as jnp

def train(w, xs, ys, lr):
    def step(w, batch):
        x, y = batch
        loss = lambda w: jnp.mean((x @ w - y) ** 2)
        return w - lr * jax.grad(loss)(w), loss(w)
    return jax.lax.scan(step, w, (xs, ys))

print(jax.make_jaxpr(train)(jnp.zeros(3), jnp.ones((6, 2, 3)), jnp.ones((6, 2)), jnp.float32(0.1)))

# { lambda ; a:f32[3] b:f32[6,2,3] c:f32[6,2] d:f32[]. let
#     e:f32[3] f:f32[6] = scan[
#       _split_transpose=False
#       jaxpr={ lambda ; g:f32[] h:f32[3] i:f32[2,3] j:f32[2]. let
#           k:f32[2] = dot_general[
#             dimension_numbers=(([1], [0]), ([], []))
#             preferred_element_type=float32
#           ] i h
#           l:f32[2] = sub k j
#           m:f32[2] = integer_pow[y=2] l
#           n:f32[2] = integer_pow[y=1] l
#           o:f32[2] = mul 2.0 n
#           p:f32[] = reduce_sum[axes=(0,)] m
#           _:f32[] = div p 2.0
#           q:f32[] = div 1.0 2.0
#           r:f32[2] = broadcast_in_dim[
#             broadcast_dimensions=()
#             shape=(2,)
#             sharding=None
#           ] q
#           s:f32[2] = mul r o
#           t:f32[3] = dot_general[
#             dimension_numbers=(([0], [0]), ([], []))
#             preferred_element_type=float32
#           ] s i
#           u:f32[3] = mul g t
#           v:f32[3] = sub h u
#           w:f32[2] = dot_general[
#             dimension_numbers=(([1], [0]), ([], []))
#             preferred_element_type=float32
#           ] i h
#           x:f32[2] = sub w j
#           y:f32[2] = integer_pow[y=2] x
#           z:f32[] = reduce_sum[axes=(0,)] y
#           ba:f32[] = div z 2.0
#         in (v, ba) }
#       length=6
#       linear=(False, False, False, False)
#       num_carry=1
#       num_consts=1
#       reverse=False
#       unroll=1
#     ] d a b c
#   in (e, f) }

unroll does not change the jaxpr. Ask for unroll=8 on an eight-step scan and the equation still says length=8 with a two-equation body; the only difference is the unroll param itself, carried along for the lowering to act on.

What it buys is measurable, and it is a trade rather than a win. On a 1024-step scan over a tanh body, going from unroll=1 to unroll=64 cut the run from 0.38 ms to 0.21 ms and pushed compile from 0.33 s to 1.07 s. Most of the run-time gain arrived by unroll=16, and most of the compile cost arrived after it.

The docstring is worth reading before you set it, because the units are easy to misread. The number says how many scan iterations happen inside a single iteration of the underlying loop, so unroll=16 on a 1024-step scan means 64 loop iterations of 16 bodies each, not 16 copies total.

import jax
import jax.numpy as jnp

xs = jnp.arange(8.)
for u in (1, 4, 8):
    jp = jax.make_jaxpr(lambda xs: jax.lax.scan(lambda c, x: (c + x, None), 0.0, xs, unroll=u))(xs)
    eqn = jp.jaxpr.eqns[0]
    print(u, eqn.params["length"], eqn.params["unroll"], len(eqn.params["jaxpr"].jaxpr.eqns))

# 1 8 1 2
# 4 8 4 2
# 8 8 8 2

Differentiate a scanned function and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → grows a twin. Two scan equations come back where there was one: the forward loop with reverse=False, and below it another with reverse=True and a carry of two rather than one.

The second one is the backward pass, written in the same primitive as the first. It starts from the last step and walks to the first, carrying the cotangent as state, which is why it needs a trip count known in advance and why the primitive that has one is the primitive you can differentiate.

Look at linear on the backward equation and it reads (True, True, False) where the forward one read all False. Those flags mark the operands autodiff knows enter linearly, the cotangent carries here, and they exist so the transpose rule can be applied to exactly those and no others.

The backward pass of a scan is a scan. Same primitive, other direction.

import jax
import jax.numpy as jnp

def total(w, xs):
    step = lambda c, x: (c * x + w, c)
    last, _ = jax.lax.scan(step, 1.0, xs)
    return last

print(jax.make_jaxpr(jax.grad(total))(jnp.float32(0.5), jnp.arange(1., 5.)))

# { lambda ; a:f32[] b:f32[4]. let
#     _:f32[] _:f32[4] = scan[
#       _split_transpose=False
#       jaxpr={ lambda ; c:f32[] d:f32[] e:f32[]. let
#           f:f32[] = convert_element_type[new_dtype=float32 weak_type=False] d
#           g:f32[] = mul f e
#           h:f32[] = add g c
#         in (h, d) }
#       length=4
#       linear=(False, False, False)
#       num_carry=1
#       num_consts=1
#       reverse=False
#       unroll=1
#     ] a 1.0 b
#     i:f32[] _:f32[] = scan[
#       _split_transpose=False
#       jaxpr={ lambda ; j:f32[] k:f32[] l:f32[]. let
#           m:f32[] = mul k l
#           n:f32[] = convert_element_type[new_dtype=float32 weak_type=True] m
#           o:f32[] = add_any j k
#         in (o, n) }
#       length=4
#       linear=(True, True, False)
#       num_carry=2
#       num_consts=0
#       reverse=True
#       unroll=1
#     ] 0.0 1.0 b
#   in (i,) }

A scan runs one body against one carry type, so the steps have to be the same shape of work. A stack of layers with different widths is not: their weights cannot be stacked into a leading axis at all, and the refusal comes from jnp.stack long before any scan sees it. Loops like that stay Python loops, and every layer contributes its own equations to the program.

How much that costs is measured next door rather than here. Chapter 12's lesson Three shapes of loop counts what an unrolled training loop puts in a program at four lengths, against the single equation the scanned form puts there, and it is worth reading before you commit to writing a long loop out. The count comes from the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, so it is exact and needs no compile to get.

The case for writing it out anyway is when you want the steps to differ deliberately: a warmup phase with a different body, a per-layer remat policy, a debug print you only want on step zero. Wrapping those in a cond inside a scan works and costs you both branches every step, which lesson 1 timed. Writing them out costs you trace size once.

The last chapter of this path assembles a real training run out of these pieces, and the mistake museum keeps the tracer-boolean failure that sends most people to these primitives in the first place. What this lesson leaves you with is the reading: one equation, seven params, and an operand list you can name before the loop ever runs.

import jax
import jax.numpy as jnp

widths = [(jnp.ones((4, 8)), jnp.zeros(8)), (jnp.ones((8, 3)), jnp.zeros(3))]
try:
    jnp.stack([w for w, _ in widths])
except ValueError as err:
    print(err)

# All input arrays must have the same shape.