Scan Patterns and Structured Control Flow

Inside JIT-compiled functions, Python's for, while, and if statements are evaluated at trace time, not runtime. This means they can't depend on runtime values (like array contents). JAX provides structured alternatives that compile into efficient XLA operations.

jax.lax.scan: Sequential Computation

import jax
import jax.numpy as jnp

# Scan replaces sequential loops with a compiled operation
# Pattern: carry, output = scan(fn, init_carry, inputs)

# Example: running sum (like cumsum but showing the pattern)
def running_sum(xs):
    def step(carry, x):
        new_carry = carry + x
        return new_carry, new_carry  # (new_carry, output)

    final_carry, outputs = jax.lax.scan(step, init=0.0, xs=xs)
    return outputs

xs = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
print(running_sum(xs))  # [1. 3. 6. 10. 15.]

# Example: apply N identical layers (weight sharing)
def apply_shared_layers(params, x, num_layers):
    def layer_fn(x, _):
        x = jax.nn.relu(x @ params['w'] + params['b'])
        return x, None  # carry = next input, no per-step output

    final_x, _ = jax.lax.scan(layer_fn, x, xs=None, length=num_layers)
    return final_x

jax.lax.while_loop: Dynamic Iteration

# while_loop for dynamic stopping conditions
def newton_sqrt(x, tol=1e-6):
    """Newton's method for square root — unknown number of iterations."""
    def cond_fn(state):
        guess, _ = state
        return jnp.abs(guess * guess - x) > tol

    def body_fn(state):
        guess, count = state
        guess = (guess + x / guess) / 2.0
        return guess, count + 1

    init_state = (x / 2.0, 0)
    final_guess, num_iters = jax.lax.while_loop(cond_fn, body_fn, init_state)
    return final_guess, num_iters

# Works under JIT!
sqrt_5, iters = jax.jit(newton_sqrt)(5.0)
print(f"sqrt(5) ≈ {sqrt_5:.6f} in {iters} iterations")

jax.lax.cond and jax.lax.switch: Branching

# Conditional execution under JIT
def safe_divide(x, y):
    return jax.lax.cond(
        y != 0,
        lambda: x / y,           # true branch
        lambda: jnp.float32(0),  # false branch
    )

# Multi-way switch (like a compiled if/elif/else)
def activation(x, choice):
    return jax.lax.switch(
        choice,
        [
            lambda x: jax.nn.relu(x),     # choice=0
            lambda x: jax.nn.gelu(x),     # choice=1
            lambda x: jnp.tanh(x),        # choice=2
            lambda x: jax.nn.swish(x),    # choice=3
        ],
        x,
    )