nn.scan issues

Boris feedback:

> Basic implementation of scan affects negatively my results with shampoo (graph per training step). I think it's because each layer is merged into one parameter and it might affect shampoo stats. I need to think of doing something smart: either concatenate params at scan or split them when passing to shampoo optimizer but it may also require custom setup methods (right now I use nn.compact)

@levskaya feedback/ideas:

> How the reverse pass is calculated. In a much larger model - useful. Giant models -> max out the performance. With scan over remat you can pencil out what’s happening/reason about what you need to do. You need `init`ializations to act like a vmap rather than a scan (init as a scan is less efficient). No loop dependance in the init fn - can treat it as a map, more efficient. `vmap` over `init`s, then `scan` over `apply`.
..scan changes the layout of the params in memory - ideally, prearrange in memory. How you run the model vs. how you save the model to disk.
🙂 There’s a doc somewhere in Anselm’s docs 🙂
This is context-dependent.