fori_loop and jacrev: it works ! but I thought it wasn't

[jecampagne]

Well,
Writing a demo nb and show that scan can perform jacrev but not fori_loop, I was astonished to discover that

@jit
def g(x):

  def body(i,val):
    v1, v2, v3, a = val 
    return v1+a[i], v2+a[i]**2, v3+a[i]**3, a

  val_init=(0.,0.,0.,x)
  N=x.shape[0]
  res = jax.lax.fori_loop(0,N,body,val_init)
  return jnp.array(list(res[:-1]))

y = jax.random.uniform(jax.random.PRNGKey(10),shape=(100,))

g(y) # DeviceArray([51.85117211, 35.0536869 , 26.71050342], dtype=float64)
jnp.sum((jacfwd(g))(y), axis=1) # DeviceArray([100.        , 103.70234421, 105.1610607 ], dtype=float64)
jnp.sum((jacrev(g))(y), axis=1) # DeviceArray([100.        , 103.70234421, 105.1610607 ], dtype=float64)

Well when I read the doc, it is clear that when the fori_loop lower and upper index are known a priori (eg. static variables) then jax.lax.scan is used,
but here I was expecting that "N" is not a priori known?

Am I wrong?, or fori_loop has changed (I am running 0.3.25 on Colab) ?

Is there a simple use-case to demonstrated that fori_loop jacrev is not working but scan does? Thanks

[jakevdp]

Thanks for the question! N here is known at compile time, because array shapes are static. If you had used an array value in place of N, then it would be dynamic and your function would not be reverse-mode differentiable.

[jecampagne]

So the slogan "always scan when you can!" is more a matter of efficiency more than impossibility to find a solution to get fori_loop working with JIT/jacrev ?

[jakevdp]

yes, I believe so

[jecampagne]

well may be the discussions 3850 is obsolete.

[jakevdp]

Yes, I think it's probably obsolete – it was right around then that fori_loop was changed to automatically use scan when possible (#2414)

[jecampagne]

Ok thanks now I have a clearer view.