我发现 jax 中的 vmap
在应用于多个参数时不会按预期运行。例如,考虑下面的函数:
def f1(x, y, z):
f = x[:, none, none] * z[none, none, :] + y[none, :, none]
return f
对于 x = jnp.arange(7), y = jnp.arange(5), z = jnp.arange(3)
,该函数的输出形状为 (7, 5, 3)
。但是,对于以下 vmap 版本:
@partial(vmap, in_axes=(none, 0, 0), out_axes=(1, 2))
def f2(x, y, z):
f = x*z + y
return f
它输出此错误:
ValueError: vmap got inconsistent sizes for array axes to be mapped:
* one axis had size 5: axis 0 of argument y of type int32[5];
* one axis had size 3: axis 0 of argument z of type int32[3]
有人可以解释一下这个错误背后的原因吗?
正确答案
vmap
的语义是它对一个或多个数组执行单个批处理操作。当您指定 in_axes=(none, 0, 0)
时,含义是“同时沿 y
和 z
的前导维度映射”:您看到的错误告诉您 y
和 y
的前导维度具有不同的大小,因此它们不兼容批处理。
您的函数 f1
本质上使用广播来编码三个批处理操作,因此要使用 vmap
复制该逻辑,您将需要 vmap
的三个应用程序。您可以这样表达:
@partial(vmap, in_axes=(0, None, None))
@partial(vmap, in_axes=(None, 0, None))
@partial(vmap, in_axes=(None, None, 0))
def f2(x, y, z):
f = x*z + y
return f