Why does jax.vmap not support broadcasting along the mapped axes?

[adonath]

I presume this has been discussed before but I could not find any reference, so I thought I bring it up as a question. I think in principle vmap could support broadcasting of the input arrays along the mapped axes. Consider the following minimal example:

import jax
from jax import numpy as jnp

a = jax.random.normal(random.key(9876), shape=(5, 4, 10))
b = jax.random.normal(random.key(976), shape=(1, 10, 4))

# batched matmul
print(jnp.matmul(a, b).shape)

# batched via vmap and repeating b
print(jax.vmap(jnp.dot)(a, jnp.repeat(b, 5, axis=0)).shape)

# why not directly support broadcasting along the mapped axes?
jax.vmap(jnp.dot)(a, b)

I can see that vmap and other functions transforms are JAX specific and do not necessarily need to support Numpy concepts such as broadcasting, but in practice it would still be useful I think and conceptional still be well defined. But maybe I have overlooked something?

[jakevdp]

It's an interesting idea – fundamentally though, broadcasting is a NumPy concept. It's supported in JAX within the jax.numpy layer, but not in general supported in jax.lax or other JAX APIs. Implicit broadcasting is sometimes convenient, but for tools like vmap we think it's better and less error-prone to be explicit.

If you want something similar to vmap that is broadcast-aware, you can use jnp.vectorize, which is implemented via vmap and does support broadcasting. In your example it would look something like this:

result = jnp.vectorize(jnp.dot, signature="(a,b),(b,c)->(a,c)")(a, b)
print(result.shape)  # (5, 4, 4)

[jakevdp]

I should add: if you want a broadcasting vmap functionality it's easy enough to implement yourself. Here's an internal utility where we do this:

jax/jax/experimental/sparse/util.py

Lines 60 to 77 in ae79249

	def broadcasting_vmap(fun, in_axes=0, out_axes=0):
	@functools.wraps(fun)
	def batched_fun(*args):
	args_flat, in_tree = tree_util.tree_flatten(args)
	in_axes_flat = flatten_axes("vmap in_axes", in_tree, in_axes, kws=False)
	size = max(arg.shape[i] for arg, i in safe_zip(args_flat, in_axes_flat) if i is not None)
	if size > 1:
	if any(i is not None and arg.shape[i] not in (1, size)
	for arg, i in safe_zip(args_flat, in_axes_flat)):
	raise ValueError("broadcasting_vmap: mismatched input shapes")
	args_flat, in_axes_flat = zip(*(
	(arg, None) if i is None else (lax.squeeze(arg, (i,)), None) if arg.shape[i] == 1 else (arg, i)
	for arg, i in zip(args_flat, in_axes_flat)
	))
	new_args = tree_util.tree_unflatten(in_tree, args_flat)
	new_in_axes = tree_util.tree_unflatten(in_tree, in_axes_flat)
	return vmap(fun, in_axes=new_in_axes, out_axes=out_axes)(*new_args)
	return batched_fun

I wouldn't recommend importing this, as it's not part of the public API, and may be removed without warning. But it's about 20 lines and you can certainly define it locally if you'd like to use it.

[adonath]

Thanks @jakevdp! I was not aware of jnp.vectorize as a convenience / Numpy compatibility layer on top of vmap. Given that both exists I fully agree that the more general explicit behavior of vmap is very reasonable and does not need further convenience. I guess it makes sense to crosslink or have a short note in the vmap docstring to mention jnp.vectorize as well.