Three shapes of loop

The Python loop from lesson one calls a compiled step two hundred times. Move that same body into jax.lax.scan and the loop stops being Python: the state becomes the carry, the step indices become the scanned input, and the per-step losses come back stacked into one (200,) array.

The two paths are not merely close. Every loss matches, and every leaf of the final state matches, as array_equal over the whole tree. That is worth confirming once yourself rather than assuming, because it is the licence to move a loop into scan after it already works.

Chapter 6 states the compile-time property that makes this attractive: scan traces its body once whatever the length, while a Python loop written inside jit unrolls. Its lesson arc prices the two against each other, in equations and in compile seconds, at three lengths. The rest of this lesson is what a scan takes away from a training run and how the run gets it back.

def body(state, step):
    return step_fn(state, data, jax.random.fold_in(base, step))


@jax.jit
def run(state, steps):
    return jax.lax.scan(body, state, steps)


def python_loop(state, n):
    losses = []
    for step in range(n):
        state, loss = train_step(state, data, jax.random.fold_in(base, step))
        losses.append(loss)
    return state, jnp.stack(losses)


start = init(jax.random.key(0))
looped, l_loop = python_loop(start, 200)
scanned, l_scan = run(start, jnp.arange(200))

print(l_scan.shape, f"{l_scan[-1]:.6f}")
print(jnp.array_equal(l_loop, l_scan))
print(all(jnp.array_equal(a, b) for a, b in zip(jax.tree.leaves(looped), jax.tree.leaves(scanned))))

# (200,) 0.005003
# True
# True

The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → says exactly what happened. One step of this model is 124 equations. The scanned version of two hundred of them is a single equation, whose primitive is scan and whose length parameter reads 200, with 126 equations in the body it carries.

126 rather than 124 because the scanned body derives its own key from the step index, so it carries a random_fold_in and a convert_element_type that the standalone step was handed ready-made. The step count lives in a parameter, not in the program's size, which is the difference between a loop the compiler sees as a loop and a loop it sees as a very long straight line.

Reading the outer jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → is also how you catch the mistake of scanning a function that is already jitted. Wrap run in make_jaxpr and the whole thing collapses to one pjit equation, because the trace stops at the jit boundary. Trace the unjitted function to see the scan.

one = jax.make_jaxpr(step_fn)(start, data, jax.random.fold_in(base, 0))
many = jax.make_jaxpr(lambda st, xs: jax.lax.scan(body, st, xs))(start, jnp.arange(200))

print(len(one.eqns))
print(len(many.eqns), many.eqns[0].primitive, many.eqns[0].params["length"])
print(len(many.eqns[0].params["jaxpr"].eqns))

# 124
# 1 scan 200
# 126

The third shape is the one people reach for when they want the whole run inside a single jit: a Python for in the traced function. It is correct, and what it costs is compile time, because the program the compiler receives carries one copy of the body per step. Chapter 6's arc prices that against a scan in equations and in seconds, and those counts are the ones to trust; they are exact where a stopwatch on a laptop is not.

What this loop adds is what the compiler prints while that happens. Compiling the hundred-step unrolled module here tripped XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s own slow-operation alarm, which prints the module name, offers to help you file a bug, and reports the duration when the compile finally lands. A training script that has ever printed those asterisks unrolled something.

Read the duration as XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s own output and not as a benchmark. Load average on this laptop was above 500 while it ran, and the claim being made is only that XLA's threshold for calling a compile slow was crossed by a program you could write by accident.

@jax.jit
def unrolled(state):
    losses = []
    for step in range(100):
        state, loss = body(state, step)
        losses.append(loss)
    return state, jnp.stack(losses)


jax.block_until_ready(unrolled(start))

# nothing on stdout, and this on stderr while it compiles:
# 2026-08-15 06:23:03.016556: E external/xla/xla/service/slow_operation_alarm.cc:73]
# ********************************
# [Compiling module jit_unrolled] Very slow compile? If you want to file a bug, run
# with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
# ********************************
# 2026-08-15 06:24:05.368570: E external/xla/xla/service/slow_operation_alarm.cc:140]
# The operation took 3m2.369747s

What scan takes away is every opportunity the host had between steps. No printing a loss as it arrives, no writing a checkpoint at step 137, no early stop, no adjusting anything from Python. The device runs all two hundred steps and comes back when it is done.

The shape that keeps both is a scan inside a Python loop. Scan fifty steps, return to the host, do whatever the host is for, scan the next fifty. Four chunks of fifty produce the same losses and the same final state as one scan of two hundred, leaf for leaf.

The chunk length is part of the signature, so the four chunks share one executable and the cache holds two entries: one for the length-200 call, one for the length-50 calls. Pick a chunk length and keep it, and the compile happens once no matter how long the run goes.

Scan the chunk, checkpoint at the boundary, scan the next.

CHUNK = 50
state = init(jax.random.key(0))
chunks = []
for c in range(4):
    state, losses = run(state, jnp.arange(c * CHUNK, (c + 1) * CHUNK))
    chunks.append(losses)            # a checkpoint can be written here

print(jnp.array_equal(jnp.concatenate(chunks), l_scan))
print(all(jnp.array_equal(a, b) for a, b in zip(jax.tree.leaves(state), jax.tree.leaves(scanned))))
print(run._cache_size())

# True
# True
# 2

The choice is not about speed first. It is about where the host stands. A Python loop over a jitted step leaves the host between every pair of steps, which is what you want while a loop is still being debugged and what makes per-step logging free to write.

A scan removes the host entirely and is the right end state for a step whose body has stopped changing. Unrolling inside one jit gives you neither the host nor a small program, and the alarm above is the reason it is a last resort rather than a default on this loop, with chapter 6's arc supplying the counts behind the general case.

Measuring which of the first two is faster on your own hardware is a job for the harness chapter 11 builds: warm up, loop, block once at the end. The gap between them is dispatch overhead against a per-step compute cost, so it depends on your model size, and a number measured on this eight-by-thirty-two MLP would tell you nothing about yours.