When to use scan and when to use vmap?

[samskiter]

I'm getting a bit confused between scan and vmap - what are the tradeoffs between them and when to use each?

Does scan with no carry effectively reduce to vmap?

How is it typically recommended to approach massing multiple training datapoints to a function? scan or vmap?

[jakevdp]

Hi - thanks for the question! The main difference is that scan always executes the body function sequentially, while vmap always executes all batches simultaneously. If you have a function for which the output of one step depends on the output of the previous step, vmap is generally not helpful.

If you have a function for which each batch is computed independently, vmap is typically the best option. The only downside of vmap is that it requires the entire computation to fit into memory at once; if vmap causes memory issues, then scan can be a good solution because it will execute steps sequentially.

Does scan with no carry effectively reduce to vmap?

No, scan will never reduce to vmap.