Try lax.scan over multiple steps.

[Drenderer]

Test what performance benefit can be gained by jax.lax.scaning over multiple training steps.
If it is significant, consider adding that functionality to klax.