Vectorizing a function that takes multidimensional input over a multidimensional array in JAX

Question

I have been trying to vectorize a function that takes two 2D arrays and return a 2D array of the same shape, so that I can apply it element wise to two 4D arrays. Here is an example:

import jax.numpy as jnp
from jax import map

M = jnp.arange(0, 400).reshape((10, 10, 2, 2))
N = jnp.arange(400, 800).reshape((10, 10, 2, 2))

def func(A, B):
    return jnp.dot(A, B)

R = vmap(func, in_axes=(1, 1))(M, N)
print(R.shape) #(10, 10, 2, 10, 2)

func is the function that I would like to apply to the 2x2 matrices that are contained in M and N. I expected the result to be of the shape (10, 10, 2, 2) because I thought that vmap applies the function to each subarray of the specified axis. Clearly I am not correctly understanding how it works. I appreciate any help! Thanks!

Answer 1

vmap will map over a single axis at a time. Since you want to map over two axes in each array, you should use two vmap calls:

R = vmap(vmap(func))(M, N)
print(R.shape) #(10, 10, 2, 2)

The reason a single vmap returns the shape it does is because you're mapping over a single axis of each array, so that the inputs to func effectively have shapes (10, 2, 2) and (10, 2, 2). If you call jnp.dot on arrays of these shapes, it will return an array of shape (10, 2, 10, 2) (the jnp.dot documentation makes clear why this is); adding the leading mapped dimension of size (10,) results in the shape you saw originally.

Answer 2

You could use jax.numpy.einsum too. This would allow you to cooperate with changing dimensions of M and N. Here is an example of 5 dimensional M and N:

import jax.numpy as jnp
from jax import vmap
from string import ascii_lowercase as letters

M = jnp.arange(0, 4000).reshape((10, 10, 10, 2, 2))
N = jnp.arange(4000, 8000).reshape((10, 10, 10, 2, 2))

C = letters[:(M.ndim - 2)] # You will map over first M.ndim-2 dimensions
subscripts = C+'XY,'+C+'YZ->'+C+'XZ' # Here I create the einsum script for the product

R1 = jnp.einsum(subscripts, M, N)

You could do the same thing with vmap. But you need three vmap calls this time:

def func(A, B):
    return jnp.dot(A, B)

R2 = vmap(vmap(vmap(func)))(M, N)

assert jnp.allclose(R1, R2)