when to use vmap

[inversecrime]

Hi,

I've been using JAX for the past few months and I'm incredibly impressed with its capabilities. Now, after gaining some experience, I would like to know about your opinion on a few matters, and I also have a suggestion/feature request.

Question:

Is it considered to be a good practice to use vmap and avoid manual vectorization whenever possible? Are there any downsides to using vmap compared to manual vectorization?

Suggestion:

One downside of vectorization is the significant memory usage. In cases where a function uses a lot of memory, it might be favorable to use conventional for loops instead - especially when each evaluation of the function takes so much time that python overhead can be neglected.

Now, let's say a function has inputs AND outputs of small size, but uses huge arrays in the computation process. Conceptually, it would make sense to use a transformation similar to vmap, since both batched input and batched output would easily fit into memory.
Such a function could use a for loop internally - maybe lax.map - and would be a huge quality-of-life improvement.

Has there been any discussion about implementing such a transformation?
I know about lax.map, but it is not as general as vmap. It would be really nice to have a transformation that could be used interchangeably with vmap - this would also allow to modify existing code without much effort.

[jakevdp]

Hi - thanks for the question and suggestion! To address some things in particular:

Is it considered to be a good practice to use vmap and avoid manual vectorization whenever possible?

No, not necessarily, although using vmap can sometimes lead to more understandable code. Performance-wise, automatically vectorized code written with vmap often lowers to an identical or near-identical sequence of XLA operations.

One downside of vectorization is the significant memory usage. In cases where a function uses a lot of memory, it might be favorable to use conventional for loops instead

This is true: whether your code is manually or automatically vectorized, it will often rely on fitting the full mapped dataset (or intermediate values) into memory: in this case, switching to scan (or lax.map/fori_loop, which are just different ways of spelling scan) can be advantageous.

Conceptually, it would make sense to use a transformation similar to vmap, since both batched input and batched output would easily fit into memory.

This is a very frequent request, and we're working on it: see #11319

[ASEM000]

Such a function could use a for loop internally - maybe lax.map - and would be a huge quality-of-life improvement.

I think jax.custom_batching.sequential_vmap is using lax.map under the hood. I am not sure if its public api yet though

[inversecrime]

Thanks for both answers!

As far as I understood, y = vmap(sequential_vmap(f))(x) will essentially do the same thing as y = vmap(f)(x), but the computation is done with lax.map instead of "real" vectorization.

Of course, a version of vmap with "chunking", as proposed in #11319, would be more general, but for now, sequential_vmap does exactly what i want.