Numpy 将 3D 数组的每个切片相乘以进行转置并求和

Question

我有一个形状为 (N, 3, 3) 的 3D numpy 数组。

我想将每个 3x3 矩阵乘以其转置，然后将它们全部相加。

这是我的代码：

m = np.zeros([3, 3])


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

但我认为只需使用

numpy

就可以完成。

我想使用

np.apply_along_axis

，但它仅适用于一维数组。

有没有办法避免

for

循环，并且一次完成乘积和求和？

Answer 1

您可以使用

numpy.einsum

：

np.einsum('ijk,ijl->kl', matrices, matrices)

示例：

np.random.seed(0)
N = 5
matrices = np.random.random((N, 3, 3))

np.einsum('ijk,ijl->kl', matrices, matrices)

array([[5.1367768 , 4.33942585, 5.02664464],
       [4.33942585, 5.73518514, 4.72942748],
       [5.02664464, 4.72942748, 6.56237687]])

循环输出：

array([[5.1367768 , 4.33942585, 5.02664464],
       [4.33942585, 5.73518514, 4.72942748],
       [5.02664464, 4.72942748, 6.56237687]])