jax.lax.scan for while loops.

[gerdm]

Hi, I'm wondering if there are any plans (or if it's feasible?) to implement jax.lax.scan for while loops. So far, I've been using the jax.experimental module and it works great, but it would be helpful to have an equivalent to scan in cases where we only need an initial condition and a function that gets called over and over until a condition is met.

[jakevdp]

Have you seen jax.lax.while_loop?

[gerdm]

Hi @jakevdp,

I have. What I'm looking for is to store intermediate values like scan does.

[jakevdp]

I don't think that would be feasible in JAX currently, because the number of iterations in a while loop is indeterminate and all JAX array sizes must be known at compilation time.

[NeilGirdhar]

You can synthesize something like what you want by doing a while-loop of scans. You'll need an iteration number in your carry for when the while loop condition was been achieved to (1) stop updating the carry, and (2) tell you how much of the output should be concatenated. You can then use concatenate the scan results from iterations of the external while loop.