Skip to content

← JAX Vectorization & Control Flow step 14 of 25

Medium Primitives

fori vs scan vs while

Jump to the task

Choose a loop by what controls it

JAX provides three structured loop primitives, each with a different API and use-case profile:

  • lax.fori_loop: an indexed loop from a lower bound to an exclusive upper bound. Body: (i, carry) → carry. Bounds can be dynamic, although this changes reverse-mode differentiation support.
  • lax.scan: a fixed-length traversal of stacked inputs, with optional stacked outputs. Body: (carry, x) → (carry, y). Emit None for y when only the final carry is needed.
  • lax.while_loop: repeat while a condition holds, useful when the iteration count depends on data. Body: state → state, with a separate condition returning a scalar boolean.

All three require the carried state's structure, shapes and dtypes to stay fixed. scan supports reverse-mode differentiation; fori_loop does when its bounds give a static trip count. A dynamically bounded fori_loop and while_loop do not support reverse-mode differentiation.

Worked mini-example

Run these snippets in order. Each sums the integers 0, 1, 2, 3, giving a final total of 6. Only the scan example also collects the intermediate totals.

Indexed loop

from jax import lax
import jax.numpy as jnp

n = 4

def add_index(i, acc):
    return acc + i

total = lax.fori_loop(
    0, n, add_index, 0
)
# total = 6

Scan with history

def add_input(acc, x):
    new = acc + x
    return new, new

carry, ys = lax.scan(
    add_input, 0,
    jnp.arange(n)
)
# carry = 6
# ys = [0, 1, 3, 6]

Condition-controlled loop

def keep_going(state):
    i, acc = state
    return i < n

def advance(state):
    i, acc = state
    return i+1, acc+i

count, total_w = lax.while_loop(keep_going, advance, (0, 0))
# count = 4
# total_w = 6

The while-loop state explicitly holds the counter and accumulator. Its final counter is 4, even though the last added index was 3. The exercise below changes the arithmetic to repeated multiplication.

Scan does not require a dummy array when there is no input sequence. Pass xs=None, length=n to run n steps; the body receives None as its second argument. Returning None as the output is a separate choice that avoids collecting a history.

Common pitfalls

  • Mixing up body signatures. fori_loop passes the index first; scan passes the carry first. while_loop passes one state value, which may itself be a tuple. Match the signature to the primitive.
  • Dynamic scan length. The length must be known when tracing, whether it comes from an input's leading axis or an explicit length. Calling int(n) on a tracer does not make it static. If jitting this exercise, mark n static at the outer boundary.
  • Changing state dtype. Use an integer counter for indexing and keep its dtype consistent. A floating-point counter can still produce a boolean condition; it is not inherently invalid for while_loop.

Problem

Implement power_three_ways(x, n) using one loop of each kind. Start each accumulator at 1 and multiply it by x exactly n times. Return a one-dimensional array of shape (3,) in this order: [fori_result, scan_result, while_result].

The base x is a real floating-point scalar. The exponent n is a non-negative whole number supplied as a scalar, possibly a float such as 3.0. Convert it with int(n) before constructing the loops. Fractional or negative exponents are outside the exercise.

For x=-2 and n=3, return [-8, -8, -8]. For n=0, return [1, 1, 1]: no multiplication occurs. This includes x=0 under the empty-product convention used here. Do not use x**n as a substitute for the loops. Output tests check values, not which primitive produced each one.

To JIT the entire function, use jax.jit(power_three_ways, static_argnames=("n",)) and pass n as a Python scalar. Calling int on a traced value does not make it static. Changing a static n can require a new compilation. The scan length must be static even though the other two loops can support dynamic bounds in other implementations.