如何用Numpy高效计算数组中多对向量的张量积？

问题描述 #

给定两个二维数组，每行对应一个向量，需要对两个数组中位置匹配的向量对分别计算张量积，再将结果组合为新数组。要求避免显式循环以提升效率（用于递归仿真场景）。

示例输入：

import numpy as np
a = np.array([[3,5,1],[6,5,4]])
b = np.array([[8,2,3],[7,8,9]])

期望输出等价于逐对计算张量积后堆叠的结果：

np.array([np.tensordot(a[0], b[0], axes=0), np.tensordot(a[1], b[1], axes=0)])

解决方案 #

以下两种方法均基于Numpy向量化操作，无Python层面循环，效率拉满：

方法1：使用np.einsum（直观且高效） #

通过指定维度运算规则，直接实现逐对张量积：

result = np.einsum('ni,nj->nij', a, b)

方法2：维度扩展+广播乘法 #

通过扩展数组维度触发广播机制，实现逐向量外积：

# 给a扩展最后一维，给b扩展倒数第二维
result = a[:, :, np.newaxis] * b[:, np.newaxis, :]

结果验证 #

以示例输入测试：

• a[0]与b[0]的张量积结果为：

  [[24,  6,  9],
   [40, 10, 15],
   [ 8,  2,  3]]
  

两种方法输出的result[0]均与上述结果一致，result[1]对应a[1]与b[1]的张量积，完全符合预期。

原理说明 #

• np.einsum('ni,nj->nij')：针对每个行索引n，将a的i维度与b的j维度做外积，最终输出n×i×j形状的数组，完美匹配逐对张量积需求。
• 维度扩展法：将a从(2,3)扩展为(2,3,1)，b从(2,3)扩展为(2,1,3)，广播后相乘会自动对每行向量执行外积运算，得到(2,3,3)的目标数组。

内容的提问来源于stack exchange，提问作者LePéler