如何用Numpy无需分配矩阵计算行存储向量数组的点积？

给定形状均为(M, 2)的数组a和b，两者各包含M个2维行向量，需求是计算对应位置向量的点积（即沿轴0的ab^T结果）。

原实现代码 #

np.diag( ( a @ b[:, :, None]).reshape(a.shape[0], a.shape[0]) )

这个实现虽然可行，但会生成一个M×M的中间矩阵，内存开销大，存在更高效的无中间矩阵实现方式：

优化方案 #

• 元素级乘法求和：直接对两个数组对应元素相乘，再沿轴1求和，完全避免中间大矩阵的生成：

  (a * b).sum(axis=1)
  

• 爱因斯坦求和约定：用np.einsum直观表述对应元素相乘后求和的逻辑，同样无中间矩阵：

  np.einsum('ij,ij->i', a, b)
  

方案说明 #

原实现通过构造高维数组做矩阵乘法，生成了不必要的M×M外积矩阵，再取对角线得到结果。优化后的两种方法都是直接针对每个行向量做计算，时间和空间复杂度均为O(M)，远优于原方案的O(M²)内存开销。

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