JAX vmap行为解析：向量化函数中点积运算结果疑问及原理探究

Let's break down exactly what's happening here with your JAX vmap code—this is a common point of confusion when first working with vectorization in JAX!

1. Why You Get the Output Shown #

First, let's recap the key setup:

• You defined v_what = vmap(what, in_axes=(None,0,None)), which tells JAX:
  • Keep a and c as single inputs for every iteration (no vectorization over their axes)
  • Vectorize over the 0th axis of b (loop over each element along this axis and run what with each one)
• Your b is a 1D array of shape (2,), so vmap will run what twice: once with each element of b (both are 2 here)

Let's walk through a single iteration of what(a, 2, c):

1. jnp.dot(a, 2): Since a is a 1D array [1,1,3] and the second argument is a scalar, jnp.dot behaves like element-wise multiplication, producing [1*2, 1*2, 3*2] = [2,2,6]
2. Adding c=1.0 gives [2+1, 2+1, 6+1] = [3,3,7]

Since we run this same calculation twice (once for each element in b), vmap stacks the two 1D results along the vectorized axis, resulting in the 2D DeviceArray([[3., 3., 7.], [3., 3., 7.]]) you saw.

2. jnp.dot Behavior After Vectorization #

The key thing to understand is how vmap transforms the inputs to your what function, which in turn changes how jnp.dot operates:

• Without vmap, calling what(a, b, c) directly would throw a shape error: you can't compute the dot product of a (3,) array and a (2,) array (their inner dimensions don't match)
• vmap solves this by splitting b along its 0th axis into individual scalars (each of shape ()). Now each iteration of what receives a scalar b instead of the full (2,) array.
• For jnp.dot(a, scalar), JAX follows standard linear algebra rules: a scalar dot product with an array is equivalent to scaling every element of the array (element-wise multiplication). This matches NumPy's behavior for np.dot with scalars.
• After computing each result, vmap collects all the 1D outputs and stacks them into a 2D array, where the first axis corresponds to the vectorized axis of the original b input.

内容的提问来源于stack exchange，提问作者ValientProcess