CODEWITHSUNDEEP
pythonarraysnumpymatrix-multiplication
Thomas Wagenaar
Thomas Wagenaar

Reputation: 6749

Computing multiple dot products of same dimension at same time

If we compute one dot product between a 4-element vector and a 4x4-matrix, we do

import numpy as np

x = np.random.rand(4)

y = np.random.rand(4,4)

print(x.dot(y))

But what if I want to do the same thing, but for 10 different vectors and matrices at the same time (but with the same dimension)? I tried:

import numpy as np

x = np.random.rand(10,4)

y = np.random.rand(10,4,4)

print(x.dot(y))

But it did not work. I also tried other methods but I could never get the desired output dimension (10,4). What is the best way to do this without using a for loop?

Upvotes: 0

Views: 143

Answers (2)

loopy walt
loopy walt

Reputation: 958

You could use matrix-matrix multiplication.

The @ operator broadcasts over everything but the last two dimensions. We can use this behaviour by temporarily inserting an extra dimension. This transforms x into a stack of 10 1x4 matrices. As the result of the matrix multiplication has the same shape as last step we remove the middle dimension.

>>> np.allclose((x[:, None] @ y)[:, 0], np.einsum("ij,ijk->ik", x, y))
# True

This method is more or less equivalent to the einsum method but marginally faster on my machine:

>>> timeit(lambda:np.einsum("ij,ijk->ik", x, y))
3.4603776139992988
>>> timeit(lambda:(x[:,None]@y)[:,0])
2.62706138700014

Upvotes: 1

Thomas Wagenaar
Thomas Wagenaar

Reputation: 6749

This seemed to work:

import numpy as np

x = np.random.rand(10,4)

y = np.random.rand(10,4,4)

np.einsum("ij,ijk->ik", x, y)

Upvotes: 0

Related Questions