python - NumPy:一次对许多小矩阵进行点积

Question

我有一长串 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
    ])

有效，但表现不佳。一个可能的改进可能是只做一个大dot产品，但我在为它正确设置A时遇到了麻烦。

有什么提示吗？

Answer 1

您可以获得转置矩阵列表作为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

@ operator只是 np.matmul 的语法糖，关于它 documentation说“如果任一参数是 N-D，N > 2，它被视为驻留在最后两个索引中的一堆矩阵并相应地广播。”