JAX vmap 行为

问题描述

我试图理解 JAX vmap 的行为，所以我写了以下代码：

import jax.numpy as jnp
from jax import vmap

def what(a,b,c):
  z = jnp.dot(a,b)
  return z + c

v_what = vmap(what,in_axes=(None,None))

a = jnp.array([1,1,3])
b = jnp.array([2,2])
c = 1.0

v_what(a,c)

输出为：

DeviceArray([[3.,3.,7.],[3.,7.]],dtype=float32)

我知道唯一被更改的输入是 b，但是有人可以解释为什么会出现这种结果吗？我对函数进行矢量化后，点积的表现如何？

解决方法

您已指定变换后的函数应映射到 b 的第一个轴上，而不是映射到 a 或 c 的任何轴上。粗略地说，您已经创建了一个执行此操作的映射函数：

def v_what(a,b,c):
  return jnp.stack([what(a,b_i,c) for b_i in b],axis=0)

对于您的输入，在每一行中，点积看起来像 jnp.dot(a,2)，结果相当于 a * 2。