Need a loop whose bound comes from an input while keeping JAX tracing and compilation intact? lax.fori_loop expresses a fixed-range loop as a structured control-flow primitive. section 3.4.1 section 5.2.1
Its form is fori_loop(lower, upper, body_fun, init_val): each iteration calls body_fun(i, val), returning the next carried value. This avoids the TracerIntegerConversionError raised by range(x) during tracing. section 5.2.1
Summing a dynamic range
def f5(x): return jax.lax.fori_loop(0, x, lambda i, v: v + i, 0) f5(5) # Array(10, dtype=int32, weak_type=True) The same function succeeds under jax.make_jaxpr() and jax.jit(); its traced representation uses a while primitive rather than an unrolled Python loop. section 5.2.1
Listing 5.16 Replacing for loop with a structured control flow primitive
def f5(x):
return jax.lax.fori_loop(0, x, lambda i,v: v+i, 0) #1
f5(5)
>>> Array(10, dtype=int32, weak_type=True) #2
jax.make_jaxpr(f5)(5) #3
>>> { lambda ; a:i32[]. let
>>> _:i32[] _:i32[] b:i32[] = while[
>>> body_jaxpr={ lambda ; c:i32[] d:i32[] e:i32[]. let
>>> f:i32[] = add c 1
>>> g:i32[] = add e c
>>> in (f, d, g) }
>>> body_nconsts=0
>>> cond_jaxpr={ lambda ; h:i32[] i:i32[] j:i32[]. let
>>> k:bool[] = lt h i
>>> in (k,) }
>>> cond_nconsts=0
>>> ] 0 a 0
>>> in (b,) }
jax.jit(f5)(5) #4
>>> Array(10, dtype=int32, weak_type=True)
Browse the complete replacement of an input-dependent Python for loop with lax.fori_loop, including traced output. section 5.2.1
