Search code examples
pythonpython-3.xnumpyarray-broadcastingnumpy-einsum

Matrix-Vector product using einsum or matmul?

I need to multiply a lot of vectors beta with the same matrix M.

Let say that the matrix M has the shape (150,7), and that the beta-s are stored in a variable of shape (7,128,128).

How would you compute the product M*beta for every element of beta?

Until know I'm doing like that:

import numpy as np

M=np.ones((150,7))
beta=np.ones((7,128,128))

result=M@(beta.reshape((7,128*128))) # the result has shape (150,128*128)
result=np.reshape(result,(150,128,128))

I'm guessing that np.einsum() could be useful here, but I don't understand how to tell it on which dimension doing the multiplication/addition.

Solution

  • Here's how you could do this using np.einsum:

    np.einsum('ij,jkl->ikl', M, beta)

    result=M@(beta.reshape((7,128*128))) # the result has shape (150,128*128)
result=np.reshape(result,(150,128,128))

np.allclose(np.einsum('ij,jkl->ikl', M, beta), result)
# True