Maximum number of iteration in lax.while_loop

[mberaha]

Is there a way to impose a maximum number of iterations in a lax.while_loop inside a jitted function?

I tried carrying on the iteration number and adding an "and" in the condition function, but I get an Abstract Tracer Value error.

For instance

from jax import jit, lax

@jit
def afunc():
    def loop_carry_fn(args):
        return args[0] + 1, args[1] * 2

    def loop_cond(args):
        return (args[0] < 1000) and (args[2] < 1e6)
 
    
    n_iter, x = lax.while_loop(
        loop_cond,
        loop_carry_fn,
        (0, 1.0))
    return x, n_iter


afunc()

[jakevdp]

You should use & instead of and:

def loop_cond(args):
  return (args[0] < 1000) & (args[1] < 1e6)

The reason is that Python's and eagerly casts its inputs to boolean, and this behavior cannot be overloaded. Thus JAX follows NumPy in using bitwise logical operators, &, |, etc. for doing element-wise logic on boolean arrays.

[mberaha]

Thanks! That was a sneaky one