NumPy：同时计算许多小矩阵的点积

我有一长串 3×3 矩阵，例如，

import numpy as np

A = np.random.rand(25, 3, 3)

对于每个小矩阵，我想执行一个外积dot(a, a.T)。列表理解

import numpy as np

B = np.array([
    np.dot(a, a.T) for a in A
    ])

有效，但表现不佳。一个可能的改进可能是做one big dot产品，但我在设置时遇到问题A正确地为它。

有什么提示吗？

您可以获得转置矩阵列表：A.swapaxes(1, 2)，以及您想要的产品列表A @ A.swapaxes(1, 2).

import numpy as np

A = np.random.rand(25, 3, 3)

B = np.array([
    np.dot(a, a.T) for a in A
    ])

C = A @ A.swapaxes(1, 2)

(B==C).all()     # => True

The @操作员只是语法糖np.matmul，关于其中文档说“如果任一参数为 N 维（N > 2），则将其视为驻留在最后两个索引中的矩阵堆栈并相应地进行广播。"