Numpy中如何用np.dot实现指定np.einsum张量缩并操作

报错原因 #

你编写的循环代码报错是因为b[:, :, k]的形状为(3, 2000)，和形状为(2000, 3)的a做np.dot运算后，输出的是形状为(2000, 2000)的矩阵，无法直接赋值给out[:, k]这个长度为2000的一维数组，触发形状不匹配错误。

正确实现 #

1. 修正后的循环版本（逻辑简单） #

如果可以接受轻量循环，修正遍历维度即可运行：

out = np.zeros((2000, 2))
for i in range(2000):
    out[i] = np.dot(a[i], b[:, i, :])

该版本遍历2000次，计算量很小，在普通设备上几乎无性能压力。

2. 纯np.dot向量化版本（无循环，效率最高） #

不需要遍历任何维度，全向量化实现：

out = np.dot(a, b.reshape(3, -1)).reshape(2000, 2000, 2)[np.arange(2000), np.arange(2000), :]

实现逻辑说明：

• 先将形状为(3, 2000, 2)的b重塑为(3, 4000)，和a做矩阵乘法得到(2000, 4000)的结果
• 将结果重塑为(2000, 2000, 2)，前两个维度分别对应a的i轴和b的i轴
• 索引取两个i轴取值相同的位置，得到最终(2000, 2)的输出

3. 可选精简实现（使用矩阵乘法运算符） #

如果允许用@（等价于np.matmul），代码更简洁易读：

out = (a[:, np.newaxis, :] @ b.transpose(1, 0, 2)).squeeze()

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