Search code examples
numpybroadcast

Broadcasting with dot product

I'm trying to solve a set of dot products. Currently, the best way I found is using this loop.

for i in range(n):
    u[:,:,i] = np.dot(K[i,:,:], v[:,:,i])

Could someone help me out to solve this by broadcasting or another faster method?

Thanks very much

Solution

  • A simple workaround is:

    u = np.matmul(K, v.transpose(2, 0, 1))
u = u.transpose(1, 2, 0)

    I am assuming K has shape n x b x c and v has shape c x d x n. Transposing v like that makes it of shape n x c x d, which makes it possible to perform n matrix multiplications of input shapes b x c and c x d and the result is n x b x d.

    Finally, you also need to transpose u to get it in the desired shape of b x d x n.