CODEWITHSUNDEEP
pythonnumpyreshape
Timofey Chernikov
Timofey Chernikov

Reputation: 55

How to multiply N vectors by N matrices using numpy?

I have a matrix M of shape (N, L) and a 3D tensor P of shape (N, L, K). I want to get matrix V of shape (N, K) where V[i] = M[i] @ P[i]. I can do it with for loop but that's inefficient, I want to do it with a single or few operations so that it would run in parallel on CUDA.

I tried just multiplying it like so

V = M @ P

but that results in a 3D tensor where V[i, j] = M[j] @ P[i].

np.diagonal(M @ P).T is basically what I want, but calculating it like that wastes a lot of computation.

Upvotes: 0

Views: 119

Answers (1)

Ivan
Ivan

Reputation: 40678

You could use np.einsum:

>>> M = np.random.rand(5, 2)
>>> P = np.random.rand(5, 2, 3)

>>> V = np.einsum('nl,nlk->nk', M, P)

>>> V.shape
(5, 3)

Upvotes: 1

Related Questions