Do functional arguments to lax.fori_loop, lax.while_loop etc. have to be pure functions?

[hr0nix]

Is this legitimate JAX code?

@jax.jit
def outer_func(a):
   def loop_body(state, idx):
       state += a
       return state
   return jax.lax.fori_loop(0, 10, loop_body, 0.0)

Here loop_body is not a pure function as it depends on the value of a in the outer scope, but the function being jitted, outer_func, is still pure.

I'm wondering whether this is a legitimate usage of lax.fori_loop, and, if not, whether I should pass every dependency of loop_body through loop state.

[jakevdp]

In general yes, functions used in higher-order primitives need to be pure in the sense of not having any side-effects. However, JAX is less strict about functions being pure in the sense of depending on external variables, even when those variables are defined within closures. The JIT-tracing machinery is able to capture loop_body's dependence on values like a in your example, and operate correctly. You should see this if you run your function while passing various values for a (on that point, I think your loop function should be defined as def loop_body(idx, state) rather than def loop_body(state, idx)).