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 _ijk =Σ _m 创建AB数组（A _ijm * B _imk ）总和仅在m指数上（重复）而不在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，也可能值得考虑使用for的显式dot循环。 BLAS矩阵乘法可能比更多解释的Python丢失的开销节省更多时间。我认为np.matmul和@应该利用dot所做的相同的事情，但我不认为np.einsum会这样做。

Answer 2

ABijk=∑m (Aijm*Bimk)转换为

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

我认为matmul运算符也会处理此

AB = A @ B

因为在最后2个维度上需要正常dot，所以将其余部分作为免费行李运送。

测试这些并让我知道它们是否有效。