Control flow

A Python if branching on a traced value raises TracerBoolConversionError, and the reason follows straight from chapter 2: a tracer carries a shape and a dtype, not a value, and if needs an actual boolean to decide which branch to take. There's nothing to decide with. The error is JAX being honest about a question it genuinely cannot answer at trace time.

It's easy to read that error as JAX simply not supporting if statements and go looking for a workaround everywhere. It's narrower than that. An if branching on a static value, something fixed at trace time like an array's shape, a static_argnums argument, or a plain Python constant closed over from outside, works exactly as written and costs nothing extra. The branch resolves once, during tracing, and only the taken branch's ops end up in the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → at all.

Python loops over a static range have the same shape of story: they run, and every one of them works, but they don't stay a loop. for i in range(1000): x = step(x) unrolls, pasting a thousand copies of step's ops into the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → one after another. The program is correct, and the compile time is proportional to the length you wrote: fine at ten iterations, a real cost at a thousand.

Fast control flow is control flow the trace can see.

When the branch really does depend on a traced value, jax.lax.cond(pred, true_fn, false_fn, *operands) is the primitive built for it, and it works by a trick worth knowing up front: both branches get staged into the compiled program, and the choice between their two results happens at run time, not trace time. Nothing about cond avoids computing the branch you don't take at the program level. What it avoids is Python ever needing to know, while tracing, which branch that will turn out to be.

For a purely elementwise choice, jnp.where is the more direct tool: no branch functions to write, just a boolean array selecting between two already-computed arrays. And it's worth knowing that under vmap, cond collapses into exactly that: since different elements of a batch might each want a different branch, the batched form has to compute both and select, so cond inside a vmapped function costs the same as where would have cost there in the first place.

import jax
import jax.numpy as jnp

@jax.jit
def flip_negatives(x):
    return jax.lax.cond(x.sum() < 0, lambda v: -v, lambda v: v, x)

print(flip_negatives(jnp.array([-3.0, 1.0])))   # [ 3. -1.]

jax.lax.scan(f, init, xs) is the primitive that makes iteration cheap to compile. f takes a carry and one slice of xs and returns a new carry plus one output; scan runs it over every slice and hands back the final carry along with every per-step output, stacked into one array. The body only gets traced once, no matter how many steps xs has, so compile time stays flat as the loop gets longer, the exact opposite of the unrolled Python loop from the first section.

The three loop primitives split cleanly on one property: whether they support reverse-mode differentiation. scan does, because it saves the residuals it needs at every step, the same bookkeeping any reverse-mode computation requires. jax.lax.while_loop does not: its trip count depends on a runtime condition, and reverse mode has no way to size a backward pass whose length isn't known in advance. jax.lax.fori_loop sits between the two: with static bounds it lowers to a scan-like form, and it falls back to while_loop's form the moment the bounds themselves become dynamic.

That split decides which primitive to reach for. Anything that iterates and carries state forward, a training loop, an RNN, a sampler, belongs in scan. A loop that has to run until some runtime condition is met, with no gradient required through it, is what while_loop is for.

import jax
import jax.numpy as jnp

def ema_step(mean, x):
    new = 0.9 * mean + 0.1 * x
    return new, new          # (carry out, output for this step)

final, path = jax.lax.scan(ema_step, 0.0, jnp.arange(5.0))
print(final)   # 0.9754...
print(path)    # every intermediate mean, stacked