NumPy中(n,n,M)与(n,n)矩阵的高效乘法实现咨询

绝对有更高效的实现方式！你的循环写法虽然直观，但没用到NumPy的向量化运算能力——这正是NumPy处理大规模数值计算的核心优势。循环会让Python解释器逐次处理每个M维度的切片，而向量化操作会把整个计算交给底层的C优化代码执行，效率提升特别明显，尤其是当M数值较大时。

下面是几种推荐的高效实现方法：

方法1：用np.einsum（最直观易读） #

np.einsum通过爱因斯坦求和约定直接描述张量运算，完美适配这种批量矩阵乘法场景：

import numpy as np

# 假设A是(n,n,M)，B是(n,n)
AB = np.einsum('ijm, jk -> ikm', A, B)

• 解释：ijm对应A的三个维度（行、列、批量索引），jk对应B的两个维度（行、列），ikm指定输出维度——对每个批量m，计算A的ij切片与B的jk矩阵相乘，得到ik结果，最终输出形状为(n,n,M)，和你的循环结果完全一致。

方法2：用np.matmul结合维度转置 #

np.matmul（或@运算符）支持广播，但需要调整维度匹配批量运算要求：

# 先把A的批量维度移到最前面，变成(M,n,n)
# 和B(n,n)相乘后得到(M,n,n)，再转置回(n,n,M)
AB = (A.transpose(2, 0, 1) @ B).transpose(1, 2, 0)

• 解释：transpose(2,0,1)把A的形状从(n,n,M)转为(M,n,n)，此时每个M对应的切片是(n,n)矩阵，和B直接做矩阵乘法后得到(M,n,n)的结果，最后再通过transpose(1,2,0)把维度转回到(n,n,M)。

方法3：用np.tensordot #

tensordot专门用于张量点积运算，通过指定轴的对应关系实现批量乘法：

# 指定A的第1个轴（列）和B的第0个轴（行）做点积
AB = np.tensordot(A, B, axes=([1], [0])).transpose(0, 2, 1)

• 解释：tensordot计算后会得到形状为(n,M,n)的结果，再通过transpose(0,2,1)调整为(n,n,M)的目标形状。

性能对比 #

这三种方法的性能都远优于你原来的循环写法。以n=100，M=1000为例，向量化方法的运行速度通常是循环写法的50~100倍（具体取决于硬件和NumPy的优化配置）。

如果你的场景中M非常大，还可以考虑结合numba库对循环进行JIT编译，但上面的NumPy原生向量化方法已经足够应对绝大多数场景了。

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