python - Numpy ndarray 乘法

Question

我有两个 3D numpy ndarray

A=np.array([[[1, 1],
             [1, 1],
             [1, 1]],

            [[2, 2],
             [2, 2],
             [2, 2]]])

B=np.array([[[ 2,  0],
             [ 0,  2]],

            [[ 2, -2],
             [-2,  2]]])

我想用元素创建 AB 数组 AB_ijk=∑_m (A_ijm*B_imk) 其中总和仅超过 m-index(重复)而不超过 i(依次重复)。

换句话说，我可以用这个 for 循环获得 di AB ndarray

for i in range(2):
    AB[i,:,:]=np.dot(A[i,:,:],B[i,:,:])

AB 等于

array([[[ 2.,  2.],
    [ 2.,  2.],
    [ 2.,  2.]],

   [[ 0.,  0.],
    [ 0.,  0.],
    [ 0.,  0.]]])

有没有办法避免 for 循环？如何用 tensordot 或 einsum 获取 AB 数组？

谢谢你的回答，我真的很感激。

Answer 1

在足够新的 NumPy (1.10+) 上，你可以做

AB = np.matmul(A, B)

或者(如果您还有 Python 3.5+):

AB = A @ B

如果你没有 NumPy 1.10+，你也可以

AB = np.einsum('ijm,imk->ijk', A, B)

对于较大的 J/M/K 维度，尤其是如果您有良好的 BLAS，可能还值得考虑使用 dot 的显式 for 循环。 BLAS 矩阵乘法可能比更多解释 Python 损失的开销节省更多时间。我认为 np.matmul 和 @ 应该利用 dot 做的相同事情，但我不认为 np .einsum 是。