Why does jax.lax.scan say the carry must have equal types?

Because the carry is the loop state and the body is traced exactly once, so its input type and its output type have to be the same type or the loop does not typecheck. The check is a function named _check_carry_type and it runs twice in scan: once to catch a wrong number of carry leaves, and once at the end on the final avals. It has two arms. If the pytree structure differs at all (a tuple where the input was an array, a dict that grew a key) you get "scan body function carry input and carry output must have the same pytree structure, but they differ", with the differing components named by path. If the structures match but any leaf fails core.typematch, you get "scan body function carry input and carry output must have equal types, but they differ", and each bad leaf is reported as having type float32[] where the output has type float32[2], followed by a clause saying whether it was the shapes, the dtypes or both. The rule is stricter than NumPy on purpose, and the docstring says why: the carry must hold a fixed shape and dtype across all iterations, and not merely be consistent up to broadcasting and promotion. In practice the two mistakes that produce this are growing the carry (concatenating onto it every step) and letting a dtype drift, usually by mixing an integer counter with a float accumulator.

Why can I not take a gradient through jax.lax.while_loop?

Because there is no transpose rule for it, and reverse-mode differentiation needs one. Forward mode works fine: while_p has a JVP rule, so jax.jvp through a while_loop runs the loop once carrying primals and tangents together. Reverse mode needs the linear part of that run to be transposed, and the function registered as the transpose of while_p does nothing but raise: "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 reason is not a missing feature. Reverse mode has to keep every intermediate the backward pass will read, and the number of intermediates is the number of iterations, which a while_loop does not know until it runs. The docstring states the constraint directly: while_loop is not reverse-mode differentiable because XLA computations require static bounds on memory requirements. If you need gradients through an iteration, the trip count has to be static, which means scan, or fori_loop with concrete bounds, which becomes a scan anyway.

Does jax.lax.fori_loop support reverse-mode autodiff?

Sometimes, and which way it goes is decided by two calls to core.is_concrete on the bounds. If both lower and upper are concrete at trace time, meaning Python integers or arrays whose value is available, fori_loop computes the trip count, builds a scan of that length, and reverse mode works. If either bound is a tracer, because you passed it as a jitted argument, it falls through to while_loop and reverse mode raises. Nothing warns you at the boundary; the same source line gives you a differentiable loop or a non-differentiable one depending on how the caller supplied the bounds. The docstring is explicit about the fork: if the trip count is static then the fori_loop is implemented in terms of scan and reverse-mode autodiff is supported, otherwise a while_loop is used and it is not. Two practical consequences. Marking the bound static under jit puts you back on the scan path. And unroll is only accepted on that path: pass it with dynamic bounds and you get a ValueError saying you can only use unroll if the loop bounds are statically known.

Why do both branches of jax.lax.cond run?

Both branches are traced, which is not the same as both being executed. cond calls pe.trace_to_jaxpr on true_fun and then on false_fun, unconditionally and in that order, because it has to have a jaxpr for each before it can emit an equation carrying both. The docstring says it plainly: both branches will be traced in all cases. So any Python side effect in a branch body, a print or an append to a list or a counter you increment, happens once per branch at trace time, regardless of the predicate. What runs later on the device is one branch: the docstring notes that using cond indicates that only one of the two branches is executed, up to compiler rewrites. There are two exceptions worth knowing. Under jax.disable_jit with a concrete predicate, cond short-circuits to an ordinary Python if and only the chosen branch is called. And under vmap with a batched predicate, the batching rule stops using a branch op at all: every branch is run on the whole batch and the results are combined with lax.select_n, so both branches really do execute.

What does the unroll argument to jax.lax.scan actually do?

It decides how many iterations of your body are inlined into one iteration of the loop XLA sees, and it is applied when the scan primitive is lowered rather than when it is traced. The default, unroll=1, gives you one rolled loop: the body appears once in the HLO. unroll=True (or unroll=0) inlines the whole thing, so a length-8 scan becomes eight copies of the body and no loop at all. An integer k in between divides the length by k: the implementation takes num_trips and remainder from divmod(length, unroll), reshapes the scanned inputs to (num_trips, k, ...), runs k inlined steps inside each iteration of a loop of num_trips iterations, and handles whatever is left over as a separate inline pass at the end. The tradeoff is the usual one for unrolling and it is entirely about the compiler: more inlined copies give XLA more to fuse across and fewer loop-carried dependencies, at the cost of a bigger module and a longer compile. Note that the carry rule is unchanged by any of this, since unroll never lets the carry change type, and that jax.lax.map takes the same argument for the same reason.

Control Flow as Primitives: One Trace, However Many Iterations

Write a Python for loop inside jit and JAX runs it at trace time, once per iteration, leaving ten thousand equations behind for a ten-thousand-step loop. The four functions in jax.lax that avoid that all work the same way: they trace the body exactly once, hang the resulting jaxpr on a single equation as a parameter, and hand XLA one loop op. The price is a rule you meet the first time you write one, and it is stricter than Python: the value carried from one iteration to the next must have the same pytree structure, the same shape and the same dtype every time round, with one narrow exception for Python scalars that costs a second trace. This page reads the real check and its error message out of the source at a pinned commit and hands you a carry-type checker you can edit, then follows the consequences: why both branches of a cond always run, why fori_loop is quietly two different primitives depending on how you passed the bounds, why scan can be differentiated in reverse and while_loop refuses outright, what happens to a branch under vmap, and what the unroll dial actually reshapes.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Control Flow as Primitives: One Trace, However Many Iterations — built to be understood, not skimmed.

Questions

Why does jax.lax.scan say the carry must have equal types?: Because the carry is the loop state and the body is traced exactly once, so its input type and its output type have to be the same type or the loop does not typecheck. The check is a function named _check_carry_type and it runs twice in scan: once to catch a wrong number of carry leaves, and once at the end on the final avals. It has two arms. If the pytree structure differs at all (a tuple where the input was an array, a dict that grew a key) you get "scan body function carry input and carry output must have the same pytree structure, but they differ", with the differing components named by path. If the structures match but any leaf fails core.typematch, you get "scan body function carry input and carry output must have equal types, but they differ", and each bad leaf is reported as having type float32[] where the output has type float32[2], followed by a clause saying whether it was the shapes, the dtypes or both. The rule is stricter than NumPy on purpose, and the docstring says why: the carry must hold a fixed shape and dtype across all iterations, and not merely be consistent up to broadcasting and promotion. In practice the two mistakes that produce this are growing the carry (concatenating onto it every step) and letting a dtype drift, usually by mixing an integer counter with a float accumulator.
Why can I not take a gradient through jax.lax.while_loop?: Because there is no transpose rule for it, and reverse-mode differentiation needs one. Forward mode works fine: while_p has a JVP rule, so jax.jvp through a while_loop runs the loop once carrying primals and tangents together. Reverse mode needs the linear part of that run to be transposed, and the function registered as the transpose of while_p does nothing but raise: "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 reason is not a missing feature. Reverse mode has to keep every intermediate the backward pass will read, and the number of intermediates is the number of iterations, which a while_loop does not know until it runs. The docstring states the constraint directly: while_loop is not reverse-mode differentiable because XLA computations require static bounds on memory requirements. If you need gradients through an iteration, the trip count has to be static, which means scan, or fori_loop with concrete bounds, which becomes a scan anyway.
Does jax.lax.fori_loop support reverse-mode autodiff?: Sometimes, and which way it goes is decided by two calls to core.is_concrete on the bounds. If both lower and upper are concrete at trace time, meaning Python integers or arrays whose value is available, fori_loop computes the trip count, builds a scan of that length, and reverse mode works. If either bound is a tracer, because you passed it as a jitted argument, it falls through to while_loop and reverse mode raises. Nothing warns you at the boundary; the same source line gives you a differentiable loop or a non-differentiable one depending on how the caller supplied the bounds. The docstring is explicit about the fork: if the trip count is static then the fori_loop is implemented in terms of scan and reverse-mode autodiff is supported, otherwise a while_loop is used and it is not. Two practical consequences. Marking the bound static under jit puts you back on the scan path. And unroll is only accepted on that path: pass it with dynamic bounds and you get a ValueError saying you can only use unroll if the loop bounds are statically known.
Why do both branches of jax.lax.cond run?: Both branches are traced, which is not the same as both being executed. cond calls pe.trace_to_jaxpr on true_fun and then on false_fun, unconditionally and in that order, because it has to have a jaxpr for each before it can emit an equation carrying both. The docstring says it plainly: both branches will be traced in all cases. So any Python side effect in a branch body, a print or an append to a list or a counter you increment, happens once per branch at trace time, regardless of the predicate. What runs later on the device is one branch: the docstring notes that using cond indicates that only one of the two branches is executed, up to compiler rewrites. There are two exceptions worth knowing. Under jax.disable_jit with a concrete predicate, cond short-circuits to an ordinary Python if and only the chosen branch is called. And under vmap with a batched predicate, the batching rule stops using a branch op at all: every branch is run on the whole batch and the results are combined with lax.select_n, so both branches really do execute.
What does the unroll argument to jax.lax.scan actually do?: It decides how many iterations of your body are inlined into one iteration of the loop XLA sees, and it is applied when the scan primitive is lowered rather than when it is traced. The default, unroll=1, gives you one rolled loop: the body appears once in the HLO. unroll=True (or unroll=0) inlines the whole thing, so a length-8 scan becomes eight copies of the body and no loop at all. An integer k in between divides the length by k: the implementation takes num_trips and remainder from divmod(length, unroll), reshapes the scanned inputs to (num_trips, k, ...), runs k inlined steps inside each iteration of a loop of num_trips iterations, and handles whatever is left over as a separate inline pass at the end. The tradeoff is the usual one for unrolling and it is entirely about the compiler: more inlined copies give XLA more to fuse across and fewer loop-carried dependencies, at the cost of a bigger module and a longer compile. Note that the carry rule is unchanged by any of this, since unroll never lets the carry change type, and that jax.lax.map takes the same argument for the same reason.

Control Flow as Primitives: One Trace, However Many Iterations

Questions

Related explainers