static_argnums and JAX control flows

[jecampagne]

Hi,
When one explicitly use jit as decorator of a function, he/she has the possibility to set some arguments as static.
But, when dealing with jax.lax.scan(body,init,xs) (idem for other control flows methods). the compilation of the body function is done w/o any control. Is there a way to pass that certain element of init are static?

[jakevdp]

There's no way to specify static variables here because scanned-over arguments cannot be static by definition (they change from iteration to iteration). The only way to have static variables in the scan body function is if they don't participate in the scan, and you can generally do this via closure, for example:

import jax
from functools import partial

def f(c, x, flag):  # flag must be static
  if flag:
    return c + 1, x
  else:
    return c, x

jax.lax.scan(partial(f, flag=True), 0, None, 10)

[jakevdp]

It’s hard for me to know how to adapt this pseudocode. I’ll reiterate the main principle: if you want loop endpoints to be static, they cannot be derived from dynamic quantities like the scan carry. Hope that makes sense

[jakevdp]

If you can provide a minimal working example (I.e. full code that executes), I might be able to give more pointed advice.

[jecampagne]

Hi @jakevdp , thanks fro your help. I understand that pseudo-code is not clear and I had already some trouble to get it... Well the originla code is not too long and is available in Colab HERE you can see that

1. there is a peice of original code (Code done with scan/while/fori loops) that is working (ie does not crash) and provide some outputs when you execute img_arr = test() and see the img_arr array.
2. the code that I am trying to get running with scan-only (and cond) is the one that is under the section Transform lax.fori_loop and lax.while in favour of lax.scan & lax.cond , you will see that it crashes for the reason you give me of course. So, if you have an idea to get it working, it would be very help full.

[jakevdp]

Thanks for that context. Looking at the full code, I don't think it's possible to write it in terms of static loops, because the size of the inner loop changes from iteration to iteration. Namely, k = jax.lax.min(m-i, mwrap-2) and i is updated each loop, and so the size of the inner loop is not static.

Your options would be either:

1. use forward-mode autodiff rather than reverse-mode autodiff with your dynamic loops
2. re-express your computation so that the size of the inner loop does not change from iteration to iteration

Hope that helps!

[jecampagne]

Ok thanks for your expertise. I will discuss with my colleagues a long your advise.