如何计算(n,1,3)与(n,3)型NumPy数组的点积得到(n,n)结果

问题需求 #

计算形状为(n,1,3)的NumPy数组（vector1）中每个子数组，与形状为(n,3)的NumPy数组（vector2）中所有子数组的点积，最终得到形状为(n,n)的结果。

尝试过的方法及问题 #

• 使用np.einsum时报错，代码如下：

  np.einsum('aij,aj->aa', vector1, vector2)
  

  错误原因：输出维度标识重复使用了相同字母a，einsum要求每个输出维度对应唯一的输入维度标识，'aa'不符合规则。

• 尝试广播+求和的方式也报错，代码如下：

  np.sum(vector1[:,np.newaxis] * vector2, axis=axis)
  

  错误原因：axis变量未定义，且广播后的形状匹配逻辑有误。

测试用例 #

vector1 = np.arange(12).reshape((4,1,3))
vector2 = np.arange(12).reshape((4,3))

正确实现方式 #

方法1：使用np.einsum #

通过调整维度标识，避免重复，即可得到(n,n)的结果：

np.einsum('ijk,bk->ib', vector1, vector2)

• 维度标识说明：
  • ijk对应vector1的(n,1,3)：i代表n维度，j代表中间的1维度，k代表3维度；
  • bk对应vector2的(n,3)：b代表n维度，k代表3维度；
  • ->ib表示输出维度对应vector1的n维度（i）和vector2的n维度（b），最终输出形状为(n,n)。

方法2：广播+求和（修正版） #

修正广播的形状匹配，指定正确的求和轴：

np.sum(vector1[:, np.newaxis] * vector2, axis=-1)

• 说明：vector1[:, np.newaxis]将vector1形状变为(n,1,1,3)，vector2自动广播为(1,n,3)，相乘后形状为(n,n,1,3)，沿最后一维（axis=-1）求和后得到(n,n,1)，若需要去掉多余的维度，可追加.squeeze()得到(n,n)。

方法3：使用np.matmul #

利用矩阵乘法的广播特性：

np.matmul(vector1, vector2.T).squeeze()

• 说明：vector2转置后形状为(3,n)，与vector1（(n,1,3)）做矩阵乘法，得到(n,1,n)，通过squeeze()去掉中间的1维度，最终得到(n,n)的结果。

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