In and out, or it does not run
A loop primitive builds its body once, against the types of the initial carry. For that single body to be correct on every iteration, whatever comes out of it has to have exactly the types that went in. Nothing adapts between steps, because there is nothing between steps: it is one traced body, run again.
One function enforces that, and both primitives call it. _check_carry_type runs at line 303 for a scan body and at line 1357 for a while_loop body, which is why the two errors read identically apart from the name in front. Learn one message and you can read both.
The message text is worth knowing by sight, because the useful half is the middle. The wrapper sentences are constant; the lines between them name the component and give you both types.
# loops.py:303, inside scan
_check_carry_type('scan body', f, init, out_tree_children[0], carry_avals_out)
# loops.py:1357, inside while_loop
_check_carry_type('while_loop body', body_fun, new_init_val, body_tree,
body_jaxpr.out_avals)
# loops.py:411-417, the raise both of them reach
raise TypeError(
f"{name} function carry input and carry output must have equal types "
"(e.g. shapes and dtypes of arrays), "
"but they differ:\n\n"
f"{differences}\n"
"Revise the function so that all output types (e.g. shapes "
"and dtypes) match the corresponding input types.")One defect, three names in front
Grow the carry and the message says shapes. A carry that starts as float32[1] and comes back as float32[2] is the shape of every accumulate-into-a-growing-buffer idea people bring from NumPy, and a loop primitive refuses it before it runs a single step.
Hand the identical defect to while_loop and the sentence that comes back is the same sentence, with a different name in front of it. That is the one check function speaking from its two call sites, which is what makes the message worth learning once instead of three times.
fori_loop is the case that repays a second look. The name in front reads scan body for a call you never wrote as a scan, and the component it names is loop_carry[1] rather than a parameter, because the loop index is carry component 0 and your value sits behind it. Both of those are the lowering showing through, and the bounds that chose that lowering get their own section later in this lesson.
The other two ways a carry breaks are told elsewhere, from the side that owns them. A dtype that drifts is read from the state side in chapter 9's lesson The step that returns its state, where the message names the offending field by its path into the state tree. A structure that changes is read from the tree side in chapter 7's lesson Structure, then leaves, together with the rule that structure is checked before any type is.
import jax
import jax.numpy as jnp
grow = lambda c: jnp.concatenate([c, jnp.ones(1)])
def show(call):
try:
call()
except TypeError as err:
print(err)
print("---")
show(lambda: jax.lax.scan(lambda c, x: (grow(c), None), jnp.zeros(1), jnp.arange(4.)))
show(lambda: jax.lax.while_loop(lambda c: c.sum() < 3.0, grow, jnp.zeros(1)))
show(lambda: jax.lax.fori_loop(0, 4, lambda i, c: grow(c), jnp.zeros(1)))
# 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 c has type float32[1] but the corresponding output carry component has type float32[2], so the shapes do not match.
#
# Revise the function so that all output types (e.g. shapes and dtypes) match the corresponding input types.
# ---
# while_loop body function carry input and carry output must have equal types (e.g. shapes and dtypes of arrays), but they differ:
#
# The input carry c has type float32[1] but the corresponding output carry component has type float32[2], so the shapes do not match.
#
# Revise the function so that all output types (e.g. shapes and dtypes) match the corresponding input types.
# ---
# 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 loop_carry[1] has type float32[1] but the corresponding output carry component has type float32[2], so the shapes do not match.
#
# Revise the function so that all output types (e.g. shapes and dtypes) match the corresponding input types.
# ---| what you called | the name in front | the component it names |
|---|---|---|
| jax.lax.scan | scan body function | c, the body parameter |
| jax.lax.while_loop | while_loop body function | c, the body parameter |
| jax.lax.fori_loop, static bounds | scan body function | loop_carry[1], one slot past the index |
A python scalar traces the body twice
Start a scan with a plain 0 and put a counter in the body. The body traces twice. The first pass traces against int32[] weak, discovers the carry comes back float32[], promotes the init, and traces again. The literal in the final jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → is 0.0, not 0.
The source says so directly, in a comment sitting above the two calls that do it. A weakly typed init may be compatible with the output despite not matching it, and the way that gets resolved is two passes, the second one with modified init values.
This is why a side effect inside a loop body is a bad place to keep a count. A print, an append to a list, a metric you increment: any of those fire once per trace, and the number of traces depends on how you wrote the init. Passing jnp.float32(0.0) instead of 0.0 is enough to make it one.
Current jax keeps the same mechanism and has added a note that it costs more than it looks: the comment on main now carries a TODO calling the two-pass approach expensive, exponential in scan nesting depth, and incomplete in the general case.
# The carry input and output avals must match exactly. However, we want to account for
# the case when init contains weakly-typed values (e.g. Python scalars), with avals that
# may not match the output despite being compatible by virtue of their weak type.
# To do this, we compute the jaxpr in two passes: first with the raw inputs, and if
# necessary, a second time with modified init values.
init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
new_init_flat, changed = _promote_weak_typed_inputs(init_flat, carry_avals, carry_avals_out)
if changed:
init = tree_unflatten(init_tree, new_init_flat)
init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
# >>> traces = []
# >>> def step(total, x):
# ... traces.append(jax.eval_shape(lambda v: v, total))
# ... return total + x, None
# >>> print(jax.make_jaxpr(lambda xs: jax.lax.scan(step, 0, xs))(jnp.arange(4.)))
# { lambda ; a:f32[4]. let
# b:f32[] = scan[
# _split_transpose=False
# jaxpr={ lambda ; c:f32[] d:f32[]. let e:f32[] = add c d in (e,) }
# length=4
# linear=(False, False)
# num_carry=1
# num_consts=0
# reverse=False
# unroll=1
# ] 0.0 a
# in (b,) }
# >>> for t in traces: print(t)
# ShapeDtypeStruct(shape=(), dtype=int32, weak_type=True)
# ShapeDtypeStruct(shape=(), dtype=float32)The loop that stops when the numbers say so
while_loop adds one requirement on top of the carry contract, and it is about the condition rather than the body. cond_fun has to return a boolean scalar. Compare an array by mistake and it hands back a bool[3], which is refused before the body is looked at.
Its trip count is decided at run time, and that decides what you can differentiate. Forward mode is fine: a tangent rides through the loop alongside the value, one step at a time, and needs no advance knowledge of how many steps there will be. The run below takes a jvp through a while_loop and gets both numbers back.
Reverse mode raises, and the message is unusually direct about the alternatives: use lax.scan, or use fori_loop with static start and stop. Both suggestions amount to the same thing, which is giving the loop a trip count known before it runs.
import jax
import jax.numpy as jnp
def grow(x):
cond_fn = lambda c: c[0] < 3.0
body_fn = lambda c: (c[0] + 1.0, c[1] * x)
return jax.lax.while_loop(cond_fn, body_fn, (0.0, 1.0))[1]
print(jax.jvp(grow, (2.0,), (1.0,)))
try:
jax.grad(grow)(2.0)
except ValueError as err:
print(err)
try:
jax.lax.while_loop(lambda v: v < 5.0, lambda v: v + 1.0, jnp.ones(3))
except TypeError as err:
print(err)
# (Array(8., dtype=float32, weak_type=True), Array(12., dtype=float32, weak_type=True))
# Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.
# cond_fun must return a boolean scalar, but got output type(s) [ShapedArray(bool[3])].Fori_loop picks its road from its bounds
Trace the same fori_loop twice with different bounds and you get two different primitives. Static bounds produce a scan with length=5 and a carry of two, the loop index riding along as the first carried value. A traced upper bound produces a while instead, with a cond_jaxpr doing the comparison and a body_jaxpr doing the work.
The trap in that is where the bound becomes traced. Outside jit, a concrete jnp.int32(5) is an ordinary value, so the scan road is taken and reverse mode works. Wrap the same call in jit and the bound is a tracer, the while road is taken, and the identical code raises the reverse-mode error.
So a gradient that works in a notebook and fails in a jitted training step is not a mystery. It is one argument that stopped being concrete, and the fix is to make the bound static rather than to restructure the loop.
import jax
import jax.numpy as jnp
decay = lambda i, c: c * 1.1 + 1.0
loop = lambda x, n: jax.lax.fori_loop(0, n, decay, x)
print(jax.make_jaxpr(lambda x: jax.lax.fori_loop(0, 5, decay, x))(jnp.float32(1.0)))
print(jax.make_jaxpr(loop)(jnp.float32(1.0), jnp.int32(5)))
print(jax.grad(loop)(jnp.float32(1.0), jnp.int32(5)))
try:
jax.jit(jax.grad(loop))(jnp.float32(1.0), jnp.int32(5))
except ValueError as err:
print(err)
# { lambda ; a:f32[]. let
# _:i32[] b:f32[] = scan[
# jaxpr={ lambda ; c:i32[] d:f32[]. let
# e:i32[] = add c 1
# f:f32[] = mul d 1.100000023841858
# g:f32[] = add f 1.0
# in (e, g) }
# length=5
# num_carry=2
# num_consts=0
# reverse=False
# ] 0 a
# in (b,) }
# { lambda ; a:f32[] b:i32[]. let
# _:i32[] _:i32[] c:f32[] = while[
# body_jaxpr={ lambda ; d:i32[] e:i32[] f:f32[]. let
# g:i32[] = add d 1
# h:f32[] = mul f 1.100000023841858
# i:f32[] = add h 1.0
# in (g, e, i) }
# body_nconsts=0
# cond_jaxpr={ lambda ; j:i32[] k:i32[] l:f32[]. let
# m:bool[] = lt j k
# in (m,) }
# cond_nconsts=0
# ] 0 b a
# in (c,) }
# 1.6105101
# Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.the same run with every scan param printed · 46 lines
import jax
import jax.numpy as jnp
decay = lambda i, c: c * 1.1 + 1.0
loop = lambda x, n: jax.lax.fori_loop(0, n, decay, x)
print(jax.make_jaxpr(lambda x: jax.lax.fori_loop(0, 5, decay, x))(jnp.float32(1.0)))
print(jax.make_jaxpr(loop)(jnp.float32(1.0), jnp.int32(5)))
print(jax.grad(loop)(jnp.float32(1.0), jnp.int32(5)))
try:
jax.jit(jax.grad(loop))(jnp.float32(1.0), jnp.int32(5))
except ValueError as err:
print(err)
# { lambda ; a:f32[]. let
# _:i32[] b:f32[] = scan[
# _split_transpose=False
# jaxpr={ lambda ; c:i32[] d:f32[]. let
# e:i32[] = add c 1
# f:f32[] = mul d 1.100000023841858
# g:f32[] = add f 1.0
# in (e, g) }
# length=5
# linear=(False, False)
# num_carry=2
# num_consts=0
# reverse=False
# unroll=1
# ] 0 a
# in (b,) }
# { lambda ; a:f32[] b:i32[]. let
# _:i32[] _:i32[] c:f32[] = while[
# body_jaxpr={ lambda ; d:i32[] e:i32[] f:f32[]. let
# g:i32[] = add d 1
# h:f32[] = mul f 1.100000023841858
# i:f32[] = add h 1.0
# in (g, e, i) }
# body_nconsts=0
# cond_jaxpr={ lambda ; j:i32[] k:i32[] l:f32[]. let
# m:bool[] = lt j k
# in (m,) }
# cond_nconsts=0
# ] 0 b a
# in (c,) }
# 1.6105101
# Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.What scan needs from xs
The length of a scan is inferred, not declared, and it comes from the leading axis of whatever you scan over. Every leaf of xs has to agree on that axis, and when they do not the error prints the sizes it found rather than a shape.
There is a form with nothing to scan over at all. Pass xs=None and a length, and you get a pure iteration: the body takes the carry and an ignored None, and the stacked output still comes back with length rows. That is the shape a sampler or a fixed-step optimizer usually wants.
Leave out both and the error is exactly what it should be. With no xs and no length, there is no number anywhere in the call that says how many times to run.
import jax
import jax.numpy as jnp
add = lambda c, x: (c + x, None)
try:
jax.lax.scan(lambda c, xy: (c + xy[0].sum(), None), 0.0, (jnp.ones((4, 2)), jnp.ones((3, 2))))
except ValueError as err:
print(err)
try:
jax.lax.scan(add, 0.0, None)
except ValueError as err:
print(err)
print(jax.lax.scan(lambda c, _: (c + 1.0, c), 0.0, None, length=3))
# scan got values with different leading axis sizes: 4, 3.
# scan got no values to scan over and `length` not provided.
# (Array(3., dtype=float32, weak_type=True), Array([0., 1., 2.], dtype=float32, weak_type=True))Check yourself
01 Your scan raised with int32[] going in and float32[] coming out. What is the smallest fix, and what would a plain 0 as the init have done instead?
Make the init match the body: jnp.float32(0.0) rather than jnp.int32(0). A plain Python 0 is weakly typed, so no error appears at all; the init is promoted and the body traces a second time against float32[].
02 A print inside a four-step scan body fired twice. What does that tell you about the init you passed?
That it held a weakly typed Python scalar whose type had to be promoted to match the carry coming out. The body traces once with the raw init and once with the promoted one, so anything with a side effect in it happens twice.
03 The same gradient runs in a notebook and raises inside jit, and the loop is a fori_loop. Why?
Because the bound was concrete outside jit, so fori_loop lowered to a scan and reverse mode worked. Inside jit the bound is a tracer, so it lowered to a while, and reverse-mode differentiation refuses a loop whose trip count is not known in advance.
Readings
- loops.py at jax-v0.4.38 ↗ scan at 129, the two-pass init promotion at 290, the carry check at 354
- jax.lax.while_loop reference ↗ the cond_fun contract, stated as the Python it stands in for
- jax.lax.fori_loop reference ↗ the docs say which lowering you get, and the differentiability that follows
- Type promotion semantics ↗ what weak means, which is what makes a python scalar init a two-pass trace