Parallelization of nested vmap

[mrtupek]

Hopefully this is an easy question:

If I have a vmap which calls a function that is turn calls another vmap, how will this be parallelized on a GPU? Will it only parallelize over the outer vmap, or will there be hierarchical parallelism over both vmaps (exposing maximum parallelism)?

[shoyer]

Yes, nested vmap will fully parallelize the entire operation.

vmap is a transformation that makes your code work exactly equivalently to manually rewriting your code with an extra dimension. So nesting vmap is equivalent to rewriting your code with an extra dimension, twice.