Matrix multiplication for multidimensional matrix

问题:

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 the new @ operator in Py3.5 +), but I'm struggling to extend this to efficiently evaluate my multidimensional arrays.

As an example, let's say I have matrix A with shape (5,3,3) and B with shape (5,3). Now, I want to matrix multiple the later parts for each 5 cases: i.e. do

res[0] = np.dot(A[0], B[0])
res[1] = np.dot(A[1], B[1])
etc

I can successfully achieve this using a loop - e.g.:

A = np.random.random((5,3,3))
B = np.random.random((5,3))

res = np.zeros([5,3])

for i in range(len(A)):
    res[i] = np.dot(A[i], B[i])

althoug this is slow because it uses a loop.

Is there a function / approach I could take to fully vectorize this please?

Thanks.

回答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]

回答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.

回答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.