numpy ndarray乘法 [英] Numpy ndarray multiplication

问题描述

我有两个3D numpy ndarray

I have two 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 (A _ijm * B _imk >)，其总和仅超过m-index(重复)，而不超过i(又是重复).

I want to create the AB array with elements AB_ijk=∑_m (A_ijm*B_imk) where the summation is only over the m-index (repeated) and not over i (that is in turn repeated).

换句话说，我可以使用此for循环获取di AB ndarray

In other words I can obtain di AB ndarray with this for loop

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循环?如何获得带有张量点或einsum的AB数组?

Is there way to avoid the for loop? How can I obtain the AB array with tensordot or einsum?

非常感谢您的回答.

推荐答案

在最新的NumPy(1.10+)上，您可以

On a sufficiently recent NumPy (1.10+), you can do

AB = np.matmul(A, B)

或(如果您也有Python 3.5 +):

or (if you also have Python 3.5+):

AB = A @ B

如果您没有NumPy 1.10+，则可以

If you don't have NumPy 1.10+, you can do

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

对于较大的J/M/K尺寸，尤其是如果您具有良好的BLAS，可能还需要考虑使用dot的显式for循环. BLAS矩阵乘法可能会比解释更多的Python所损失的开销节省更多的时间.我认为np.matmul和@应该利用dot所做的相同操作，但是我不认为np.einsum这样做.

For large J/M/K dimensions, especially if you have a good BLAS, it might also be worth considering the explicit for loop with dot. The BLAS matrix multiply might save more time than the overhead of more interpreted Python loses. I think np.matmul and @ are supposed to take advantage of the same things dot does, but I don't think np.einsum does.

这篇关于numpy ndarray乘法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！