nn.scan vs. lax.scan performance for LSTMCell

[minqi]

I understand the lifted version, nn.scan, allows wrapping LSTMCell in a new module LSTM, which then handles the state required for performing a scan-style rollout over a sequence. However, aside from this syntactic benefit, is there any expected performance improvement?

[cgarciae]

Hey @minqi, performance-wise there should be no differences between the two. nn.scan mainly handles collections and rngs appropriately, it also introduces some features of its own but underneath it calls JAX's scan.