Performance difference between while_loop inside vmap vs. other way

[femtomc]

Question as in title: I have two options to implement a piece of code. Initially, I started by writing a version of my code with a single replicate, with a while_loop and then vmap across the top to run multiple replicates.

However, I'm wondering if it would be a better idea to apply vmap across replicates first, and then put things in a while_loop (with a termination condition that looks across a status flag across all replicates, and decides to stop or keep going).

What's the code generation look like for these two variants? Is one always better than the other?

[jakevdp]

If I understand your question correctly, I think you'll find that the two approaches will be effectively equivalent. A vmap of a while_loop essentially creates a while_loop that doesn't terminate until all mapped loops terminate. So whether you do that automatically via vmap of while_loop or manually via while_loop of vmap you'll probably see broadly the same performance.