JAX Pathways
Core Foundation · Local CPULaptop CPULesson 3 of 4 in Phase 03

Phase 03 / State, randomness & control flow

Compiled loops with lax.scan

Revision v0.7.0 ·

Lesson · about 65 minutes · CPU practice

Replace a fixed-length Python loop with scan and compare final states.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Pytrees and structured parameters
Hardware: CPU

Download notebook
Other downloadsPython worked examplesLesson text
Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder
py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment
python phases/03-state/03-compiled-loops-with-lax-scan/code/main.py

PowerShell command · repository folder
.\.venv\Scripts\python.exe phases/03-state/03-compiled-loops-with-lax-scan/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Check: History: [1.,1.5,1.75,1.875][1., 1.5, 1.75, 1.875]; Final: 1.8751.875.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 03.03 Stage 1 of 5 · Understand
1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence
Stage 1 of 5 · Understand the MechanismCore mental model, equations & visual walkthrough

What you will be able to do

  • Separate carry from recorded outputs
  • Keep the carry contract invariant
  • Differentiate the recurrence by hand
  • Compile the repeated transition without a timing promise

The problem

Suppose a simulator keeps part of its current value and adds a new input at each step. Let’s write out four steps so we know what should happen. Then we’ll express the same loop with scan, compare it with Python, and see how changing the retained fraction affects the final value.

The idea

A scan separates the state carried between steps from the outputs collected at each step. This is useful whenever the same rule repeats over time while the state keeps a fixed structure, shape and dtype.

Carry state through a compiled loop
  1. Initial statecarry₀, inputs[0]

    Start with an explicit state tree.

  2. Bodycarry₁, output₀ = body(carry₀, input₀)

    Return the next state and this step’s output.

  3. Repeatcarryₜ → body → carryₜ₊₁

    scan collects outputs and returns the final carry.

Every carry leaf keeps the same shape and dtype through the loop. Outputs are stacked along the scan axis; the carry is the state that connects steps.

Conceptual illustration. Verify the computation in the lesson’s Python experiments.

Unroll three steps before compiling the loop

For a separate hand-worked recurrence, start with carry 00 and inputs [1,2,3][1,2,3]. Let each step add its input to the carry and emit the new carry. The successive carries are 1,3,61,3,6; the final carry is 66, while the stacked outputs are [1,3,6][1,3,6].

Those two returned objects answer different questions. Keeping only the final carry discards the intermediate outputs. Keeping outputs does not mean the carry itself grows on each iteration. A fixed carry contract is what makes the repeated computation structured.

Trace one horizontal state edge and one downward output edge in the diagram. When adapting a Python loop, write down which variables belong on each edge before translating the loop to lax.scan.

Carry moves through time; outputs are collected

Predict: What changes if the step emits the old carry instead of the new carry?

An additive scan starts with carry \(0\) and inputs \(1,2,3\). Horizontal arrows pass updated carry; downward arrows collect emitted values. Final carry \(6\) and outputs \([1,3,6]\) are different return values. Arrow lengths do not represent runtime.
Architecture and dataflow mechanism diagram.

An additive scan starts with carry 00 and inputs 1,2,31,2,3. Horizontal arrows pass updated carry; downward arrows collect emitted values. Final carry 66 and outputs [1,3,6][1,3,6] are different return values. Arrow lengths do not represent runtime.

Open full-size diagram

Separate carry from recorded outputs

Each scan step receives (carry,input) and returns (next_carry,output). Carry is what the next step needs. Output is what you want stacked into a history. They can have different meanings even when this demonstration returns the new state in both positions.

For decay 0.50.5, initial state 00 and four inputs of 11, the states are 11,1.51.5,1.751.75,1.8751.875. There are four recorded outputs; the initial state is not automatically part of history. Decide whether a plot should include that initial point before labeling the time axis.

(state_t, input_t) → transition → (state_(t+1), recorded_t)
state: 0 → 1 → 1.5 → 1.75 → 1.875

Keep the carry contract invariant

Carry values change, but their pytree structure, shapes and dtypes must remain compatible across steps. Growing a Python list or concatenating a larger array into the carry changes that contract. Allocate a fixed-size state or use the separately stacked outputs instead.

This constraint is a representation decision, not an arbitrary restriction on simulation. A fixed-dimensional physical state naturally meets it. If a system has variable numbers of entities, consider padding and masks with an explicit maximum capacity.

Differentiate the recurrence by hand

With four unit inputs and initial zero, the final value is 1+d+d2+d31+d+d^2+d^3. Its derivative is 1+2d+3d21+2d+3d^2, giving 2.752.75 at d=0.5d=0.5. The derivative also obeys a recurrence: st+1=statet+dsts_{t+1}=\mathrm{state}_t+d s_t, starting from s0=0s_0=0.

Autodiff through scan differentiates the represented state transition. It does not make an unstable or inaccurate numerical model valid. Verify a small analytic case first, then a changed input sequence, before trusting a longer simulation.

sfinal(d)=1+d+d2+d3dsfinaldd=1+2d+3d2sfinal(0.5)=1.875,sfinal′(0.5)=2.75\begin{aligned}s_{\mathrm{final}}(d)&=1+d+d^2+d^3\\\frac{ds_{\mathrm{final}}}{dd}&=1+2d+3d^2\\s_{\mathrm{final}}(0.5)&=1.875,\qquad s_{\mathrm{final}}\prime(0.5)=2.75\end{aligned}

Compile the repeated transition without a timing promise

A Python loop inside jit can be unrolled while tracing; scan expresses a repeated operation compactly. That can reduce compilation work for long loops, but it is not a universal promise that scan is faster for every workload.

Before benchmarking, compare numerical results on the same inputs and synchronize execution. Keep step count and precision explicit. For this lesson we verify correctness and derivatives; hardware profiling belongs to a later course unit.

The carry approaches a fixed point

Predict before looking: Why does the increment between successive states shrink?

The carry approaches a fixed point. The solid line shows the carry after each completed scan step. It starts at step \(1\) with value \(1\), then rises through \(1.5\), \(1.75\), and \(1.875\). The dashed horizontal line marks the fixed point \(2\). The initial carry \(0\) is not plotted. Each visible point is the value after applying the recurrence once more, so the first point already includes one update.
Plot from a recorded CPU computation · JAX 0.9.2 · 1 CPU device

Read the figure

The solid line shows the carry after each completed scan step. It starts at step 11 with value 11, then rises through 1.51.5, 1.751.75, and 1.8751.875. The dashed horizontal line marks the fixed point 22.

The initial carry 00 is not plotted. Each visible point is the value after applying the recurrence once more, so the first point already includes one update.

Connect it to the computation

The update is ct+1=0.5ct+1c_{t+1}=0.5c_t+1. Subtracting it from the fixed point gives 2−ct+1=0.5(2−ct)2-c_{t+1}=0.5(2-c_t): each step halves the remaining gap. This explains both the upward movement and its slowing pace.

The shrinking gaps are 11, 0.50.5, 0.250.25, and 0.1250.125. scan carries the evolving value forward while collecting this history. One more step would produce 1.93751.9375; it would not add another full unit or restart from the initial carry.

Open full-size figureDownload data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Stage 2 of 5 · Build & Run the ReferenceStep-by-step code construction & verified execution output

Where to put & run this code: Work in phases/03-state/03-compiled-loops-with-lax-scan/code/main.py (or downloaded exercises/state-03.py / exercises/state-03.ipynb) and execute from the repository root with python3 phases/03-state/03-compiled-loops-with-lax-scan/code/main.py.

1 Prepare the inputs

Create main.py in your lesson workspace. Add this first block; use the environment from setup.

main.py · append build step 1
# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

These explicit inputs define the case that the later checks will verify.

2 Build the computation

Append this block below the inputs in the same file.

main.py · append build step 2
# Step 2 — Build the computation: step returns the next scalar carry and one recorded scalar.
def step(carry, increment):
    next_value = 0.5 * carry + increment
    return next_value, next_value
increments = jnp.array([1., 1., 1., 1.])

Code building blocks in this step

  • jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.

step returns the next scalar carry and one recorded scalar. scan feeds that carry back into step and stacks the four recorded outputs.

3 Run and check the result

Append the checks, save main.py, and run python main.py from this folder using your course environment.

main.py · append build step 3
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
# Run compiled structured control flow via `jax.lax` (`(final, history)`).
final, history = jax.lax.scan(step, jnp.array(0.), increments)
print("History:", history)
print("Final:", float(final))
assert jnp.allclose(history, jnp.array([1.,1.5,1.75,1.875]))
assert jnp.allclose(final, history[-1])

Code building blocks in this step

  • jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
  • jax.lax.scan(step_fn, init_carry, xs, length=...) — Compiles a sequential loop where step_fn(carry, x) returns (next_carry, y), returning (final_carry, stacked_ys).

Compare the output to the expected result below before making the exercise change.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build
# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

# Step 2 — Build the computation: step returns the next scalar carry and one recorded scalar.
def step(carry, increment):
    next_value = 0.5 * carry + increment
    return next_value, next_value
increments = jnp.array([1., 1., 1., 1.])

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
# Run compiled structured control flow via `jax.lax` (`(final, history)`).
final, history = jax.lax.scan(step, jnp.array(0.), increments)
print("History:", history)
print("Final:", float(final))
assert jnp.allclose(history, jnp.array([1.,1.5,1.75,1.875]))
assert jnp.allclose(final, history[-1])

Run the example

python3 phases/03-state/03-compiled-loops-with-lax-scan/code/main.py
About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example
main.py · complete Python example
# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp
# Step 2 — Build the computation: step returns the next scalar carry and one recorded scalar.
def step(carry, increment):
    next_value = 0.5 * carry + increment
    return next_value, next_value
increments = jnp.array([1., 1., 1., 1.])
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
# Run compiled structured control flow via `jax.lax` (`(final, history)`).
final, history = jax.lax.scan(step, jnp.array(0.), increments)
print("History:", history)
print("Final:", float(final))
assert jnp.allclose(history, jnp.array([1.,1.5,1.75,1.875]))
assert jnp.allclose(final, history[-1])

Expected result: History: [1.,1.5,1.75,1.875][1., 1.5, 1.75, 1.875]; Final: 1.8751.875.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:45:30.167199+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

History: [1.    1.5   1.75  1.875]
Final: 1.875
History: [1.    1.5   1.75  1.875]
Final: 1.875
Expected carry-shape mismatch
PASS: state-03

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Stage 3 of 5 · Controlled ExperimentsPredict before running & test boundary conditions

Compare scan to the Python recurrence

Compare scan to the Python recurrence · Python experiment
# Experiment — Compare scan to the Python recurrence: The loop is an independent control-flow reference and the fixed...
inputs=jnp.array([1.,-1.,2.,0.])
state=0.
reference=[]
for inc in [1.,-1.,2.,0.]:
    state=0.5*state+inc
    reference.append(state)
# Run compiled structured control flow via `jax.lax` (`(last, observed)`).
last,observed=jax.lax.scan(step,jnp.array(0.),inputs)
assert jnp.allclose(observed,jnp.array([1.,-0.5,1.75,0.875]))
assert jnp.allclose(observed,jnp.array(reference))
assert jnp.allclose(last,state)

Expected: History [1,−0.5,1.75,0.875][1,-0.5,1.75,0.875]; final 0.8750.875.

The loop is an independent control-flow reference and the fixed numbers verify time indexing.

Verify state and sensitivity together

Verify state and sensitivity together · Python experiment
# Experiment — Verify state and sensitivity together: The scalar objective is the final state.
# Define `terminal(d)` to carry state across steps with `jax.lax.scan`:
def terminal(d):
    return jax.lax.scan(lambda c,u:(d*c+u,d*c+u),jnp.array(0.),increments)[0]
# Differentiate the objective to obtain `(value, sensitivity)` via automatic differentiation.
value,sensitivity=jax.value_and_grad(terminal)(0.5)
assert jnp.allclose(value,1.875)
assert jnp.allclose(sensitivity,2.75)
assert jnp.allclose(jax.jit(terminal)(0.5),value)

Expected: Final state 1.8751.875 and sensitivity 2.752.75.

The scalar objective is the final state. Returning the whole history to ordinary grad would need a deliberate scalar reduction.

Stage 4 of 5 · Hands-On Exercise & CheckpointModify the code yourself, solve transfer tasks & test your understanding
Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/03-state/03-compiled-loops-with-lax-scan/code/main.py (or inside the Make it yours cell of exercises/state-03.ipynb) and run python3 phases/03-state/03-compiled-loops-with-lax-scan/code/main.py to verify your prediction before opening the reference solution.

Make the decay an explicit scalar argument. Differentiate the final value with respect to decay and compare at 0.50.5 with the analytic derivative 2.752.75.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
  • jax.grad(loss_fn)(params, ...) — Transforms a scalar-output function into a function returning the gradient PyTree with the same structure as params.
  • jax.lax.scan(step_fn, init_carry, xs, length=...) — Compiles a sequential loop where step_fn(carry, x) returns (next_carry, y), returning (final_carry, stacked_ys).

2. Step-by-step implementation plan:

  1. Define simulate(decay) to carry state across steps with jax.lax.scan:
Starter code scaffold · copy into your file & fill in the TODOs 
# Exercise solution: Make the decay an explicit scalar argument.
# Define `simulate(decay)` to carry state across steps with `jax.lax.scan`:
def simulate(decay):
    def transition(carry, increment):
        value = ...  # TODO: compute value
        return ...  # TODO: return computed result
    return ...  # TODO: return computed result
assert jnp.allclose(jax.grad(simulate)(0.5), 2.75)  # TODO: complete assertion check
Reveal the reference solution
Reference solution
# Exercise solution: Make the decay an explicit scalar argument.
# Define `simulate(decay)` to carry state across steps with `jax.lax.scan`:
def simulate(decay):
    def transition(carry, increment):
        value = decay * carry + increment
        return value, value
    return jax.lax.scan(transition, jnp.array(0.), increments)[0]
assert jnp.allclose(jax.grad(simulate)(0.5), 2.75)

Record a different output

Practice

Carry the state, but record its squared value. Confirm that final carry still means state.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
  • jax.lax.scan(step_fn, init_carry, xs, length=...) — Compiles a sequential loop where step_fn(carry, x) returns (next_carry, y), returning (final_carry, stacked_ys).

2. Step-by-step implementation plan:

  1. Run compiled structured control flow via jax.lax ((last, energy)).
Starter code scaffold · copy into your file & fill in the TODOs 
# Record a different output (Practice): Carry and history need not have the same interpretation.
# Run compiled structured control flow via `jax.lax` (`(last, energy)`).
last,energy = jax.lax.scan(...)  # TODO: compute last,energy
assert jnp.allclose(last,1.875)  # TODO: complete assertion check
assert jnp.allclose(energy,jnp.array([1.,2.25,3.0625,3.515625]))  # TODO: complete assertion check
Show a hint

The second returned value is only the recorded output.

Reveal reference solution and reasoning
Record a different output · Python solution
# Record a different output (Practice): Carry and history need not have the same interpretation.
# Run compiled structured control flow via `jax.lax` (`(last, energy)`).
last,energy=jax.lax.scan(lambda c,u:(0.5*c+u,(0.5*c+u)**2),jnp.array(0.),increments)
assert jnp.allclose(last,1.875)
assert jnp.allclose(energy,jnp.array([1.,2.25,3.0625,3.515625]))

Carry and history need not have the same interpretation. The state remains unsquared for the next transition.

Diagnose a growing carry

Challenge

Try appending each input to the carry with concatenate. Explain why scan rejects it and replace it with stacked output.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
  • jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
  • jax.lax.scan(step_fn, init_carry, xs, length=...) — Compiles a sequential loop where step_fn(carry, x) returns (next_carry, y), returning (final_carry, stacked_ys).

2. Step-by-step implementation plan:

  1. Run the boundary check and catch the expected exception:
  2. Run compiled structured control flow via jax.lax ((_, saved)).
Starter code scaffold · copy into your file & fill in the TODOs 
# Diagnose a growing carry (Challenge): The error identifies incompatible carry input/output types.
def growing(c,u):
    return ...  # TODO: return computed result
# Run the boundary check and catch the expected exception:
try:
    jax.lax.scan(growing,jnp.zeros((0,)),increments)
except TypeError:
    print("Expected carry-shape mismatch")
else:
    raise AssertionError("Expected scan type failure")
# Run compiled structured control flow via `jax.lax` (`(_, saved)`).
_,saved = jax.lax.scan(...)  # TODO: compute _,saved
assert jnp.array_equal(saved,increments)  # TODO: complete assertion check
Show a hint

The carry shape is part of the loop type.

Reveal reference solution and reasoning
Diagnose a growing carry · Python solution
# Diagnose a growing carry (Challenge): The error identifies incompatible carry input/output types.
def growing(c,u):
    return jnp.concatenate([c,u[None]]),u
# Run the boundary check and catch the expected exception:
try:
    jax.lax.scan(growing,jnp.zeros((0,)),increments)
except TypeError:
    print("Expected carry-shape mismatch")
else:
    raise AssertionError("Expected scan type failure")
# Run compiled structured control flow via `jax.lax` (`(_, saved)`).
_,saved=jax.lax.scan(lambda c,u:(c,u),jnp.array(0.),increments)
assert jnp.array_equal(saved,increments)

The error identifies incompatible carry input/output types. Store per-step values in scan outputs instead of growing the carry.

Check your understanding

What must remain consistent across scan steps?

Choose one answer

Diagnose the result

The error identifies incompatible carry input/output types. Store per-step values in scan outputs instead of growing the carry.

Stage 5 of 5 · Keep Your EvidenceRecord your verification receipt & carry takeaways forward

Carry forward

  • The loop is an independent control-flow reference and the fixed numbers verify time indexing.
  • The scalar objective is the final state. Returning the whole history to ordinary grad would need a deliberate scalar reduction.

Keep your evidence

Download evidence template

Keep the labeled carry/output transition, Python-reference trajectory, polynomial derivative, compiled value check and growing-carry failure with a fixed-shape repair.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references