The step that returns its state

Trace a function that reads an array from the enclosing scope and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → has something before the semicolon that chapter 2's softmax never had. { lambda a:f32[3]; b:f32[3]. ... } names two slots: constvars on the left of the semicolon, invars on the right. The closed-over bias went left. The argument x went right.

The two slots differ in when they receive a value. An invar gets a fresh one on every call, because it is an argument. A constvar received its value once, during the trace, and the executable carries that array inside itself from then on.

Rebinding the Python name is what makes the second half visible. After bias = jnp.zeros(3) the compiled function still adds ones, because it holds the array that was there when it traced, not the name that used to point at it. Nothing raises and nothing warns. A loop that reassigns a global every step keeps computing with the value from step zero.

State you read from scope is a constvar. State you pass is an invar.

import jax
import jax.numpy as jnp

bias = jnp.ones(3)

def add_bias(x):
    return x + bias

print(jax.make_jaxpr(add_bias)(jnp.zeros(3)))

fast = jax.jit(add_bias)
print(fast(jnp.zeros(3)))
bias = jnp.zeros(3)          # rebind the name the function reads
print(fast(jnp.zeros(3)))

# { lambda a:f32[3]; b:f32[3]. let c:f32[3] = add b a in (c,) }
# [1. 1. 1.]
# [1. 1. 1.]

Thread the same value as an argument instead and it moves to the other slot, where it can be replaced. The step below carries (count, total), and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → opens with a:i32[] b:f32[3] c:f32[3]: the state tuple flattened into two invars, then the batch. It closes with (d, e, g), which is the new count, the new total, and the per-step output, in that order. One leaf of state, one invar, one outvar.

The equation list also prices the counter. f:f32[] = convert_element_type[new_dtype=float32 weak_type=False] d is in the program because an i32[] count met a float division, and the recording will not do that silently. A state field's dtype is visible work, not an implementation detail of your dataclass.

That one-leaf-one-slot correspondence is what the next lesson's alias map indexes into. Donation is expressed per output position against per input position, so a state tree with eight leaves gives the compiler eight independent decisions to make.

import jax
import jax.numpy as jnp

def step(state, x):
    count, total = state
    count = count + 1
    total = total + x
    return (count, total), total / count

print(jax.make_jaxpr(step)((jnp.int32(0), jnp.zeros(3)), jnp.ones(3)))

# { lambda ; a:i32[] b:f32[3] c:f32[3]. let
#     d:i32[] = add a 1
#     e:f32[3] = add b c
#     f:f32[] = convert_element_type[new_dtype=float32 weak_type=False] d
#     g:f32[3] = div e f
#   in (d, e, g) }

(state, x) -> (new_state, y) is not only how a training step is written. It is the exact signature jax.lax.scan wants of a body, which chapter 6 covers from the loop's side. Read from the state's side, the claim is narrower and more useful: the carry is your state tree, and the two names describe one object.

The carry adds one requirement a Python loop never enforces. What comes out has to have the same type as what went in, leaf for leaf, and count + 1.0 instead of count + 1 is enough to break it: an int32[] carry component comes back as float32[] and scan refuses before running anything.

Read the refusal closely and it names state[0], rooted at the body's own parameter name and indexed by position in the flattened tree. On a deep state that path is how you find the field, without bisecting the step by hand.

import jax
import jax.numpy as jnp

def step(state, x):
    count, total = state
    return (count + 1.0, total + x), total    # count leaves as a float

init = (jnp.int32(0), jnp.zeros(3))
try:
    jax.lax.scan(step, init, jnp.ones((4, 3)))
except TypeError as err:
    print(err)

# scan body function carry input and carry output must have equal types (e.g.
# shapes and dtypes of arrays), but they differ:
#
# The input carry component state[0] has type int32[] but the corresponding
# output carry component has type float32[], so the dtypes do not match.
#
# Revise the function so that all output types (e.g. shapes and dtypes) match
# the corresponding input types.

Run the step eight times in a Python loop, then hand the identical function to scan over the identical inputs. The final carry matches, and so does the last per-step output: (Array(8, dtype=int32), Array(36., dtype=float32)) and 3.5 from both.

Nothing about the step was edited between the two runs. The Python loop rebinds a name eight times and dispatches eight programs; scan hands the carry from one iteration to the next inside one. Which of those you want is a compile-time and memory question that chapters 6 and 11 both weigh.

The portability is the part worth keeping. A step that threads its state runs unchanged under a loop, under scan, under jit, and under grad. A step that mutates something outside itself runs correctly under none of them, and the first section showed what it computes instead.

import jax
import jax.numpy as jnp

def step(state, x):
    count, total = state
    return (count + 1, total + x), total / (count + 1)

xs = jnp.arange(1.0, 9.0)

state = (jnp.int32(0), jnp.float32(0.0))
for x in xs:
    state, mean = step(state, x)

scanned, means = jax.lax.scan(step, (jnp.int32(0), jnp.float32(0.0)), xs)
print(state, mean)
print(scanned, means[-1])
print(state[1] == scanned[1], mean == means[-1])

# (Array(8, dtype=int32), Array(36., dtype=float32)) 3.5
# (Array(8, dtype=int32), Array(36., dtype=float32)) 3.5
# True True