Skip to content

← JAX Vectorization & Control Flow step 13 of 25

Easy Primitives

x^n via lax.fori_loop

Jump to the task

Repeat an update over an index range

lax.fori_loop(lower, upper, body, initial) calls body(i, carry) for indices from lower up to, but not including, upper. Each call returns the carry for the next iteration. With lower=0 and upper=3, the indices are 0, 1 and 2. If upper is no greater than lower, no iterations run and the initial carry is returned.

The bounds can be runtime integer values; they do not have to be known at compile time. This makes fori_loop useful under JIT when a Python range(n) would require a concrete n. Static Python loops are still valid, but can be unrolled into larger compiled programs.

Use the index, or ignore it

import jax
from jax import lax

@jax.jit
def sum_n(n):
    return lax.fori_loop(0, n, lambda i, total: total + i, 0)

sum_n(5)  # 0 + 1 + 2 + 3 + 4 = 10
sum_n(3)  # 0 + 1 + 2 = 3

The index comes first in the body arguments. For a power, the body does not need the index: it just multiplies the running product by x. With x=2 and n=3, the carry goes 1 → 2 → 4 → 8. Return only the new carry from the body, not an (index, carry) pair.

Dynamic bounds and differentiation

When the trip count is known during tracing, fori_loop can use scan and support reverse-mode differentiation. With dynamic bounds, it uses a while loop and reverse-mode differentiation is not supported. Runtime bounds are valid for computing values; this limitation is about differentiating that computation.

Use integer bounds. This exercise supplies n as a scalar that may be represented as a float, so convert it with jnp.int32(n). That conversion works on a traced value without requiring a Python int. It does not make n static. The input contract excludes fractional exponents; truncating one would not compute a fractional power.

Problem

Implement fori_power(x, n) using n repeated multiplications with lax.fori_loop. The base x is a real floating-point scalar. The exponent n is a non-negative whole number within the int32 range, supplied as a scalar such as 3.0.

Convert n with jnp.int32, start the carry at 1.0, and multiply by x once per step. Return the final scalar. Do not use x**n in place of the loop. The carry must keep its shape and dtype across updates.

For x=-2 and n=3, return -8. For n=0, return 1, including when x=0: this follows the empty-product convention because no update runs. Negative or fractional exponents are outside this exercise. The output tests check results but cannot enforce which loop you use.