jax.lax.scan with "inpure" body function

[mchagneux]

Hi all,

I have a small doubt regarding part of my code when I re-read the documentation for lax.scan. To give an example I often write code that looks like this:

length = 50
def g(z):
 def f(carry, x):
  return carry + z, None
 return jax.lax.scan(f, init=0., xs=None, length=length)

i.e. the body of the scanned-over function depends on the parameter of the outer function. Now I read that "scan compiles f" so how come I don't have any problems with the fact that f captures a "global" parameter (i.e. outside of its scope ?). Is it that whenever I call g with a new argument then f is recompiled, because f is a functional attribute of g ? Then what happens if I compile the outer function (g) ? I haven't had any numerical surprises so far but I just want to check that I'm not missing anything.

[jakevdp]

This is fine because f closes over z, and z never changes without f being re-defined in a new function call. What would be problematic is if you changed z within the scope of the function and scanned over f again without re-defining it.

[mchagneux]

Ok thanks, very clear. I still have a doubt regarding compilation times: since f is redefined for every function call of g, will it be recompiled everytime I call g ? If I understand correctly, since scan compiles the function, it technically "sees" a new function f to compile every time g is called. In this case would I get better performance by putting z in the carry of the scan even though its doesn't change throughout the loop, so that JAX can store a compiled version of f. But maybe it's what happens already under the hood ? And if I compile g, I guess none of this matters anymore ?

[jakevdp]

Yes, as written f will be re-compiled each time you call g. You can fix this by defining f outside the scope of g, or by wrapping g itself with jit, in which case the body of g will only be traced once for each z of matching shape/dtype.