Skip to content

← JAX Vectorization & Control Flow step 12 of 25

Medium Primitives

lax.while_loop vs Python while

Jump to the task

Put a data-dependent condition inside the compiled loop

A Python while statement needs a concrete boolean before it can choose whether to enter its body. Under jax.jit, a condition that depends on a runtime input is traced without that input's value. Converting that condition to a Python boolean raises TracerBoolConversionError.

lax.while_loop keeps the condition and body inside a compiled loop, so the condition can be evaluated at runtime. Its state must keep the same structure, shapes and dtypes across iterations.

import jax
from jax import lax

@jax.jit
def bad(x):
    while x > 1.0:
        x = x * 0.5
    return x

@jax.jit
def good(x):
    return lax.while_loop(lambda s: s > 1.0, lambda s: s * 0.5, x)

# bad(4.0) raises TracerBoolConversionError
good(4.0)  # 1.0: 4 -> 2 -> 1, then the condition is false

A Python loop controlled entirely by static values can still work under JIT. Other transformations have their own tracing behavior; for example, grad without an enclosing JIT can sometimes differentiate through Python control flow using concrete primal values. The issue is whether Python needs to decide a condition whose value is unknown, not a blanket ban on while statements in any transformed function.

Count updates, not condition checks

In this exercise, multiply a value by 0.9 until it is strictly below a threshold. Carry (count, current_x) and continue while current_x >= threshold. Each body call increments count once.

Starting at 1 with threshold 0.5, the values after six and seven updates are approximately 0.531441 and 0.478297. Return 7. If the initial value is already below the threshold, return 0. If it equals the threshold, perform one update and return 1.

The comparison is intentionally >=. Using > would stop too early at equality. Repeated floating-point multiplication also rounds, so replacing the loop with a logarithm formula can change a count near a boundary. Count the actual updates requested here.

Problem

Implement safe_iterate_until(x, threshold) with lax.while_loop. Repeatedly update current_x = current_x * 0.9 while current_x >= threshold. Return the number of updates as a scalar float32, using jnp.float32(count) after the loop. The loop counter itself can remain an integer.

Inputs are positive finite float32 scalars. For this exercise, use thresholds at least jnp.finfo(jnp.float32).tiny, the smallest normal positive float32 value. Zero, negative and subnormal thresholds are outside the exercise: floating-point rounding or underflow can otherwise prevent a decreasing sequence from reaching the requested strict bound.

Return only the count, not the final value or a history. Output tests check counts but cannot enforce which loop produced them. As with other dynamic while loops, computing the result does not imply reverse-mode differentiation through the loop is supported.