Skip to content

← JAX Vectorization & Control Flow step 11 of 25

Hard Primitives

Bounded Search via lax.while_loop

Jump to the task

Stop at a match or a search limit

To find the first zero, keep an index as the loop state. Advance it while the current value is nonzero and the index is still inside the allowed prefix. lax.while_loop expresses this data-dependent stop under JIT, where a Python condition on a traced value cannot decide whether to break out of a loop.

Let limit = min(max_iters, len(x)). The searchable positions are 0 through limit-1. Reaching limit means no match was found in that prefix, even if a zero exists later in the array.

x max_iters Result Meaning
[1, 2, 0, 4] 10 2 First zero is at index 2
[1, 2, 0, 4] 2 2 Search stopped before index 2
[0, 3] 10 0 First element matches
[1, 2] 10 2 No zero in the array

The same returned number can mean different things. A match was found only when the returned index is less than limit. With a zero limit, return zero without advancing, regardless of x's contents.

A boolean bound does not protect an array access

This condition looks guarded, but is not safe:

(idx < limit) & (x[idx] != 0.0)

Both operands of & are evaluated. Once idx reaches len(x), the expression still indexes past the end. JAX can clamp an out-of-bounds read instead of raising an ordinary Python IndexError, so a plausible returned result does not prove the indexing was valid. An empty array has no valid element to read at all.

Handle an empty array first, using its static shape. For a nonempty array, use safe_idx = jnp.minimum(idx, len(x)-1) for the read, while keeping idx < limit as the logical search condition. At the stopping bound the condition may read a valid element again, but cannot accept it as a match or advance the index. This makes the array access valid even though & evaluates both sides.

Python and is not a replacement: it asks Python to convert the traced condition to a boolean. Use array boolean operations together with safe indexing. Use jnp.minimum for a runtime max_iters too; Python min would need a Python comparison.

Problem

Implement while_find_first_zero(x, max_iters) using lax.while_loop. The input x is a one-dimensional real floating-point JAX array. max_iters is a non-negative whole-number scalar within int32 range, possibly supplied as a float. Convert it with jnp.int32 and cap it at x.shape[0]. Negative or fractional limits are outside this exercise.

Return a scalar int32: the first index whose value equals zero within the allowed prefix, or limit if no such value exists. An empty array or zero limit returns zero. Test exact equality with zero, not whether the value is negative or merely close to zero.

Initialize the index with jnp.int32(0) to make the output dtype explicit. A plain Python zero is not inherently an invalid loop carry. Output tests check the index, but cannot enforce use of while_loop.