Matrix Multiplication For Multidimensional Matrix (/array) - How To Avoid Loop?

I'm trying to evaluate a matrix multiplication with arrays containing multiple matrices to be multiplied together. This can easily be achieved with two matrices using np.dot (or th

Solution 1:

You can use np.einsum -

np.einsum('ijk,ik->ij',A,B)

Using np.matmul -

np.matmul(A,B[...,None]).squeeze()
np.matmul(A,B[...,None])[...,0]

Solution 2:

Also, using sum and broadcasting:

np.sum(A * B[:, np.newaxis, :], axis=2)

Einsum seems to be the fastest option though. This would be the slowest since it instantiate the product of elements.

Solution 3:

Because I always have a hard time understanding einsum(), I came up with this:

np.diagonal(np.squeeze(np.dot(A, B[:,:,None])), axis2=2).T

It works, and it has no loops, but it is significantly slower than einsum() because of the :,:,None expansion of B from 2D to 3D (which is later reduced by diagonal()). I'm not sure if there's a way to make this more compact.