如何用Numpy优化三维数组乘法中的嵌套for循环？

你的核心操作是对每个维度l，计算矩阵a[:, :, l]和b[:, :, l].T的矩阵乘法——原循环里(a[j, :, l] * b[m, :, l]).sum()等价于矩阵乘法中第j行第m列的元素。以下是几种纯Numpy的高效实现方式：

方法1：使用np.einsum（语义最直观） #

einsum可以直接描述张量运算规则，完全匹配你的需求：

import numpy as np

a = np.random.rand(640, 640, 20)
b = np.random.rand(39, 640, 20)
result = np.einsum('jkl,mkl->jml', a, b)

解释：jkl对应a的维度，mkl对应b的维度，对共同的k维度求和，输出jml对应result的(640,39,20)形状，和原循环结果完全一致。

方法2：矩阵乘法+维度交换 #

通过调整数组维度，利用Numpy的批量矩阵乘法能力：

result = np.matmul(a, b.swapaxes(1,2)).swapaxes(1,2)

解释：b.swapaxes(1,2)将b转为(39,20,640)，a为(640,640,20)，matmul会自动对最后两个维度做矩阵乘法，得到(640,39,20)的结果，刚好匹配result的形状。

方法3：使用np.tensordot #

通过指定求和轴实现：

result = np.tensordot(a, b, axes=([1], [1])).transpose(0,2,1)

解释：tensordot(a,b,axes=([1],[1]))对a的第1轴和b的第1轴求和，得到(640,20,39)，再转置为(640,39,20)即可。

以上方法均基于Numpy底层C实现的向量运算，效率远高于Python手动循环。

你提到的第二个循环中，m的范围随l变化（range(-l, l+1)），对应需要计算的m索引不连续且数量动态调整。这种场景下，Numpy的批量张量运算无法直接覆盖所有情况，因为Numpy要求操作的维度是规整的。如果仅用Numpy，只能保留循环结构，无法实现类似初始场景的高效批量优化。

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