如何在NumPy中更简洁实现张量的逐空间单元矩阵-向量乘法？

嘿，这个场景我太熟悉了！你现在用嵌套循环逐个处理空间单元的矩阵-向量乘法，虽然逻辑清晰，但在高维张量场景下不仅代码冗余，效率也远不如NumPy的向量化操作。下面给你两个简洁高效的实现方式，不管是2D空间维度还是更高维的(m,n,o,p)都能完美适配：

方法一：用np.einsum——维度对应直观，适合复杂张量操作 #

np.einsum通过爱因斯坦求和约定直接描述张量运算关系，完全不需要循环。针对你的需求，一行代码就能搞定：

C = np.einsum('ab,bmn->amn', A, B)

• 索引解释：A的形状是(l,l)对应索引a,b；B的形状是(l,m,n)对应索引b,m,n；最终得到的C是(l,m,n)对应索引a,m,n。这个表达式精准对应你循环里的A@B[:,i,j]逻辑，自动对每个(m,n)位置的向量执行矩阵乘法。
• 高维场景扩展：如果B是(l,m,n,o,p)的张量，只需要扩展索引即可：

  C = np.einsum('ab,bmnop->amnop', A, B)
  

方法二：用reshape + 矩阵乘法——代码简洁，快速上手 #

如果你觉得einsum的索引有点绕，也可以通过维度重塑简化操作：

C = (A @ B.reshape(l, -1)).reshape(l, m, n)

• 核心思路：把B后面的所有空间维度（m,n）压缩成一维，让B变成(l, m*n)的矩阵，和A（(l,l)）完成矩阵乘法后，再把结果重塑回原来的空间维度(l,m,n)。
• 高维场景适配：同样不需要修改核心逻辑，直接替换目标形状即可：

  C = (A @ B.reshape(l, -1)).reshape(l, m, n, o, p)
  

验证正确性 #

你可以用原来的循环结果验证这两种方法的准确性：

import numpy as np

l,m,n = 2,3,4
A = np.random.rand(l,l)
B = np.random.rand(l,m,n)

# 原循环实现
C_loop = np.zeros((l,m,n))
for i in range(m):
    for j in range(n):
        C_loop[:,i,j] = A@B[:,i,j]

# einsum实现
C_einsum = np.einsum('ab,bmn->amn', A, B)
# reshape实现
C_reshape = (A @ B.reshape(l, -1)).reshape(l, m, n)

# 检查结果一致性
print(np.allclose(C_einsum, C_loop))  # 输出 True
print(np.allclose(C_reshape, C_loop)) # 输出 True

为什么这两种方法更好？ #

这两种都是向量化操作，NumPy会利用底层的BLAS/LAPACK优化库加速计算，比Python嵌套循环快几个数量级——尤其是当张量维度很大时，效率提升会非常明显。同时代码更简洁，后期维护也更省心。

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