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.