sum matrix elementwise using vmap (jax)?

user_5791398 0
Viewed 550

I'm trying to understand the in_axes and out_axes options in vmap. For example, I want to sum two matrix and get the output with the same shape.

X = np.arange(9).reshape(3,3)
Y = np.arange(0,-9,-1).reshape(3,3)
def sum2(x,y):
    return x + y
vmap(sum2,in_axes=((0,1),(0,1)))(X,Y)

I think I mapped both axes 0 and 1 for X and Y respectively. The output will have the same shape as X,Y. But i get the error,

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-403-103694166574> in <module>
      3 def sum2(x,y):
      4     return x + y
----> 5 vmap(sum2,in_axes=((0,1),(0,1)))(X,Y)

    [... skipping hidden 2 frame]

~/anaconda3/lib/python3.8/site-packages/jax/api_util.py in flatten_axes(name, treedef, axis_tree, kws)
    276       assert treedef_is_leaf(leaf)
    277       axis_tree, _ = axis_tree
--> 278     raise ValueError(f"{name} specification must be a tree prefix of the "
    279                      f"corresponding value, got specification {axis_tree} "
    280                      f"for value tree {treedef}.") from None

ValueError: vmap in_axes specification must be a tree prefix of the corresponding value, got specification ((0, 1), (0, 1)) for value tree PyTreeDef((*, *)).
jax python
1 Answers

First of all, the easiest way to do an element-wise sum is to use the built-in broadcasting of binary operations, and call sum2(X, Y) directly.

That said, if you're trying to understand vmap: the issue is that vmap can only map one axis at a time. If you want to map multiple axes, you can nest multiple vmaps. I believe what you intended to do can be expressed this way:

from jax import vmap
import jax.numpy as np

X = np.arange(9).reshape(3,3)
Y = np.arange(0,-9,-1).reshape(3,3)

def sum2(x,y):
    assert x.ndim == y.ndim == 0
    return x + y

vmap(vmap(sum
  vmap(sum2, in_axes=(0, 0), out_axes=0),
  in_axes=(1, 1), out_axes=1
)(X,Y)

Note: I added the assertion about number of dimensions to demonstrate that the mapped function is being called on scalar values.

Also, notice that when the mapped axes match, e.g. in_axes=(0, 0) can be equivalently written in_axes=0, but I left it as a tuple because it was closer to the syntax you were trying.

In fact, a far more concise way to do the same computation with nested vmap would be to use the default arguments: vmap(vmap(sum2))(X, Y) will do the same elementwise sum.

user_29378310
Related