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=         


        

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.