vmap across arbitrary leading dimensions.

[andrewwarrington]

I wanted to vmap a function across arguments with a number of batch dimensions that varies depending on use, i.e. one call may have no batch dimensions, whereas in another place it may have three leading batch dimensions. This idea is referenced here and #1168, although it appears to have had no resolution.

---

The first type of solution I came up with was something like this (illustrated for the case of matrix-vector multiplication)

def adaptive_matmul(x, A):

  single_matmul = lambda _x: A @ _x
  vmapped_once = jax.vmap(lambda _x: single_matmul(_x))
  vmapped_twice = jax.vmap(lambda _x: vmapped_once(_x))
  vmapped_thrice = jax.vmap(lambda _x: vmapped_twice(_x))

  if x.ndim == 1:
    y = single_matmul(x)
  elif x.ndim == 2:
    y = vmapped_once(x)
  elif x.ndim == 3:
    y = vmapped_twice(x)
  elif x.ndim == 4:
    y = vmapped_thrice(x)
  else:
    raise RuntimeError("I've run out of batch dimensions!")

  return y

However this only supports up the three batch dimensions, and so is pretty tedious and limited, and is really ugly (although it is mainly for illustrative purposes).

---

But we can express this as a recursion and get something more general like this:

def vmap_with_arbitrary_leading_dimensions(fn, leading_dims, *args):
  """
  vmap a function over arguments with arbitrary leading dimension.

  Assumes that the *args have the same leading batch dimensions, i.e. 
  `in_axes=(None,...)` is not supported.  

  :param fn:            (Callable) Function to apply.
  :param leading_dims:  (Int) The number of leading dimensions in the args.
  :param *args:         (Tuple) The arguments to the function with `leading_dims`
                        leading dimensions. 
  """
  if leading_dims == 0:
    return fn(*args)
  else:
    next_vmapped_fn = jax.vmap(vmap_with_arbitrary_leading_dimensions, 
                               in_axes=(None, None, *((0, ) * len(args))))
    return next_vmapped_fn(fn, leading_dims - 1, *args)

which we can call as follows:

# Define whatever your function is.
# Ax + c -> y
# x.dim x const.dim -> y.dim
A = jnp.ones((20, 15))                      # (y.dim, x.dim).  Fixed param.
fn = jax.jit(lambda _x, _c: (A @ _x + _c))  

# Simulate some batched data we apply the function to.
_single_x = jnp.ones((15, ))                # (x.dim, )
_batch_repeats = (5, 10, 50, 1)             # (*batch_dims, 1).

# We will assume we can at least programatically access how many leading dims in the variable there is.
leading_dims = len(_batch_repeats) - len(_single_x.shape)

# Build the variables
x = jnp.tile(_single_x, _batch_repeats)     # (*batch_dims, x.dim)
c = jnp.ones(x.shape[:leading_dims])        # (*batch_dims, c.dim)

# Apply the function.
result = vmap_with_arbitrary_leading_dimensions(fn, leading_dims, x, c)

The result is then the return of fn applied across the batch dimensions.

There are some deviations here from the JAX-ey way of doing things; namely that vmap_with_arbitrary_leading_dimensions as proposed returns the result, whereas vmap returns a function that can then be evaluated. But if fn is jitted maybe there isn't a huge performance hit? I think you require concrete instantiations of args anyway to work out how many recursions are needed, but maybe there is a way around this.

In some ways, though, I think this is more JAX-ey that my naive alternatives: in the same way that you should write fn for a single application and vmap across a "1-d" batch, you should be able to apply vmap across batches with arbitrary dimension.

---

I'd love to know if there is a better way of going about this, either that is already in JAX or otherwise, and whether this would be a good thing to add.

Thanks,
Andy

[jakevdp]

It sounds like what you want is similar to the semantics of jnp.vectorize? It might look something like this, using the variables you defined:

from functools import partial

@partial(jnp.vectorize, signature='(m),()->(n)')
def fn(x, c):
  return A @ x + c

out = fn(x, c)

This is implemented via nested calls to vmap, as in your approach. But if that's not the API you want, I think your wrapper function looks like a great solution.

[andrewwarrington]

oh. sometimes it really is that simple eh? 😅😂

Thank you for posting this! This is exactly what I was looking for. I was sort-of disbelieving that it wouldn't already be in JAX somewhere, but I didn't realise it would be that simple.

Thanks!
A

[andrewwarrington]

Also linking this Stack Overflow answer, hilariously by jakevdp as well, giving some more context.

https://stackoverflow.com/questions/69099847/jax-vectorization-vmap-and-or-numpy-vectorize