Numpy ndarray multiplication

Question

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]]])

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).

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,:,:])

and AB is equal to

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

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

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

Thank you for your answers, I really appreciate it.

Answer 1

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

AB = np.matmul(A, B)

or (if you also have Python 3.5+):

AB = A @ B

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

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

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.

Answer 2

ABijk=∑m (Aijm*Bimk) translates to

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

I think the matmul operator will also handle this

AB = A @ B

since it takes the normal dot on the the last 2 dimensions, carrying the rest along as free baggage.

Test those and let me know if they work.