Structured control flow with lax

Real algorithms loop and branch on data; lax primitives keep that code jittable and compile-fast.

When control flow depends on array values, Python if/for cannot be traced. jax.lax provides primitives that become single operations in the compiled program:

lax.cond(pred, true_fn, false_fn, *operands) picks a branch at run time. Both branches are traced and must return the same shapes and dtypes. lax.switch(index, branches, *operands) generalizes to several.
lax.fori_loop(lower, upper, body_fn, init) runs body_fn(i, carry) repeatedly.
lax.while_loop(cond_fn, body_fn, init) loops until cond_fn(carry) is false.
lax.scan(f, init, xs) walks over the leading axis of xs, threading a carry, and stacks per-step outputs: f(carry, x) -> (carry, y). It is the workhorse for RNNs, ODE solvers and cumulative computations.

Why not just unroll a Python loop? A 10,000-step unrolled loop yields a huge program and very long compile times; scan or fori_loop compiles the body once.

Differentiation rules matter:

cond, scan and fori_loop with static bounds are reverse-mode differentiable.
while_loop (and fori_loop with traced bounds) supports forward mode only; use scan with a fixed length when you need grad.

The carry must keep the same shape and dtype every iteration, so initialise it with the right dtype (e.g. 0.0 not 0 for a float accumulator).

Example

import jax
import jax.numpy as jnp
from jax import lax

@jax.jit
def safe_log(x):
    return lax.cond(x > 0, jnp.log, lambda v: jnp.zeros_like(v), x)
print(safe_log(jnp.e), safe_log(-1.0))      # 1.0 0.0

def cumsum_scan(xs):
    def step(carry, x):
        carry = carry + x
        return carry, carry                  # (new carry, per-step output)
    total, ys = lax.scan(step, 0.0, xs)
    return total, ys
print(cumsum_scan(jnp.array([1.0, 2.0, 3.0])))  # (6.0, [1. 3. 6.])

def power(x, n):
    return lax.fori_loop(0, n, lambda i, acc: acc * x, 1.0)
print(jax.jit(power, static_argnums=1)(2.0, 10))  # 1024.0

print(jax.grad(lambda x: cumsum_scan(x)[0])(jnp.ones(3)))  # [1. 1. 1.]

Remember

lax.cond/switch branch on traced values; all branches must return matching types
lax.scan threads a carry over a sequence and compiles the body once
Use scan rather than while_loop when you need reverse-mode gradients

Watch out for

Starting a float accumulator with the int 0, which makes the carry dtype mismatch and raises a TypeError.

Check yourself

You need grad through a loop of 1,000 steps whose body depends on a traced value. Which construct fits best?

Official docs: docs.jax.dev/en/latest/201/control-flow.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks