lifted nn.scan unroll seems not work comparing jax.lax.scan

[luweizheng]

Hi,

I am doing AI4Science research on solving differential equations. I mainly use scan to simulate different timesteps. I find flax's lifted transformations can be used in nn.Module. It would be more concise to write code.

I implement scan in two versions. In the pure jax function version, Module is transformed into jax function using apply. And I can use jax.lax.scan when calculating loss. In lifted version, I use nn.scan inside nn.Module. Now both of the two versions have same speed when unroll=1 is set. For pure jax function version, when I set unroll to higher value, time for a single iteration significantly reduces and time for compilation grows a little bit. But for flax's version, unroll seems does not take any effect. Running time for bigger unroll does not reduce.

I also check the source code, and it seems that in axes_scan.py unroll is passed into lax.scan.

[jheek]

The scan wrapper in linen/transforms.py passes the default (1) instead of the unroll value passed PR #2213 will fix it

[luweizheng]

@jheek
Thanks a lot! So maybe I can use the new nn.scan in next minor release?

[luweizheng]

Just test the main branch. unroll now works!