如何用Numpy实现3D数组各切片转置相乘求和（替代循环）

我有一个形状为(N, 3, 3)的3D Numpy数组，希望将每个3×3矩阵与其转置相乘，再将所有结果求和。我的实现代码如下：

import numpy as np
m = np.zeros([3, 3])

for matrix in matrices:
    m += np.dot(matrix.T, matrix)

但我认为仅用Numpy即可完成此操作。我曾尝试使用np.apply_along_axis，但它仅适用于一维数组。请问是否有方法避开for循环，一步完成乘积与求和操作？

当然可以用Numpy的广播和矩阵运算直接实现，无需手动循环，以下是几种高效实现方式：

方法一：使用np.einsum（直观灵活） #

np.einsum可以通过维度标记清晰定义运算逻辑，完美匹配需求：

import numpy as np
# matrices为形状(N,3,3)的输入数组
m = np.einsum('nij,nkj->ik', matrices, matrices)

逻辑解释：

• nij表示输入数组的维度：n是第0轴的N个矩阵样本，i和j是单个矩阵的行、列
• nkj对应每个矩阵转置后的维度（原矩阵的列j转为行k，这里k与i对应）
• 最终输出ik维度的3×3矩阵，自动完成「每个矩阵转置相乘 + 沿N轴求和」的操作

方法二：批量矩阵乘法+轴求和 #

利用转置和批量矩阵乘法实现，逻辑更贴近手动循环的思路：

# 一步式写法
m = (np.transpose(matrices, axes=(0,2,1)) @ matrices).sum(axis=0)

分步拆解：

1. np.transpose(matrices, (0,2,1))：将每个3×3矩阵转置，数组形状保持(N,3,3)
2. @运算符：执行批量矩阵乘法，每个转置后的矩阵与原矩阵相乘，得到形状(N,3,3)的中间结果
3. .sum(axis=0)：沿第0轴（N个矩阵的维度）求和，得到最终的3×3结果矩阵

方法三：使用np.tensordot #

通过指定轴的对应关系完成运算：

m = np.tensordot(matrices, matrices, axes=([0,1], [0,2]))

逻辑解释：
指定第一个数组的轴0（N）、轴1（行），与第二个数组的轴0（N）、轴2（列）对应，自动完成转置乘和求和操作。

以上方法均基于Numpy底层优化实现，比手动Python循环效率高得多，尤其适合N值较大的场景。

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