如何使用JAX的vmap实现矩阵逐元素求和并正确配置in_axes参数

• vmap 单次调用仅支持沿单个轴做向量化映射，每个输入参数对应的 in_axes 配置只能是单个整数（指定映射的轴序号）或者 None（不沿任何轴映射），不能传入轴序号组成的元组。
• 你传入的 ((0,1),(0,1)) 相当于给每个输入数组指定了两个要映射的轴，而输入的X、Y都是单个数组（属于PyTree的叶子节点），配置的参数结构和输入的PyTree结构不匹配，因此触发报错。

你要实现两个形状为(3,3)的矩阵逐元素相加，输出相同形状的结果，有两种常用配置方式：

方式1：单层vmap（适用运算支持向量输入的场景） #

如果你的自定义函数本身支持向量/张量运算（比如示例中的加法本身就支持广播和向量运算），直接指定沿0轴映射即可：

import jax.numpy as np
from jax import vmap

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

def sum2(x,y):
    return x + y

# 沿0轴做映射，每个批次输入是两个长度为3的行向量，相加后vmap自动拼接为(3,3)的输出
res = vmap(sum2, in_axes=(0, 0))(X,Y)

方式2：嵌套vmap（适用函数仅支持标量输入的场景） #

如果你的自定义函数只能处理标量输入，需要嵌套两层vmap分别映射0轴和1轴：

# 外层映射0轴，内层映射1轴，会逐元素调用sum2后拼接回原形状
res = vmap(vmap(sum2, in_axes=(1, 1)), in_axes=(0, 0))(X,Y)

两种方式输出的结果都和直接计算X + Y完全一致，满足形状要求。

内容的提问来源于stack exchange，提问作者H.Yuanchen