Best practices for nested loop

[rkruegs123]

I have a typical JAX issue. I want to implement behavior like the following:

for i in range(n):
    for j in range(i):
        # do something

The value of n is variable, so I have to use fori or scan. Clearly, the inner loop is an issue as it depends on i.

I understand there are several options for implementing this, e.g. masking. However, I was wondering if there are any best practices for implementing this behavior. For example, maybe gradients are known to propagate better using one method vs. another. Or, if there are reports of relatively poor gradient quality using masking, I would consider making n a static argument and recompiling every time.

I should add that I can't vectorize the inner loop as later values depend on previously computed values. However, this does raise a sub-question I have about fori. My understanding is that the following form is allowed:

for i in range(n):
  arr = jnp.arange(0, i)
  # vectorized computation with `arr`

If I am correct, why is this jittable but the nested for loops are not? Is their dependence on i not the same?

Any and all feedback would be much appreciated. Thank you!

[jakevdp]

Lots of questions here; some general answers:

• the most performant approach, particularly if you're running on accelerators, will be to re-express your algorithm in terms of vectorized computations wherever possible (i.e. avoiding scan or fori_loop) even if it means computing more results than strictly necessary.
• Gradients will also be more performant using vectorized operations, although they will work correctly with scan & fori_loop (not with while_loop)
• Regular loops are jit-compilable because they are unrolled during tracing – as far as the operations produced by jit, the result is identical to the case where you replace the for loop with n literal copies of the code it contains.

Hope that helps