使用vmap（jax）对矩阵元素求和？

Question

我试图理解vmap中的in_轴和out_轴选项。 例如，我想对两个矩阵求和，得到相同形状的输出

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)

我想我分别为X和Y映射了轴0和轴1。输出将具有与X，Y相同的形状。 但我得到了错误

---------------------------------------------------------------------------
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((*, *)).

Answer 1

首先，进行元素求和的最简单方法是使用内置的二进制操作广播，并直接调用sum2(X, Y)

也就是说，如果您试图理解vmap：问题是vmap一次只能映射一个轴。如果要映射多个轴，可以嵌套多个VMAP。我相信你的意图可以这样表达：

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)

注意：我添加了关于维数的断言，以证明映射函数是在标量值上调用的

另外，请注意，当映射的轴匹配时，例如in_axes=(0, 0)可以等价地写入in_axes=0，但我将其保留为元组，因为它更接近您尝试的语法

事实上，使用嵌套的vmap执行相同计算的一种更简洁的方法是使用默认参数：vmap(vmap(sum2))(X, Y)将执行相同的元素求和