CODEWITHSUNDEEP
pythonnumpylinear-algebra
Joce
Joce

Reputation: 2342

Vectorized inner product with numpy

I need to compute many inner products for vectors stored in numpy arrays.

Namely, in mathematical terms,

Vi = Σk Ai,k Bi,k

and

Ci,j = Σk Ai,k Di,j,k

I can of course use a loop, but I am wondering whether I could make use of a higher level operator like dot or matmul in a clever way that I haven't thought of. What is bugging me is that np.cross does accept arrays of vector to perform on, but dot and inner don't do what I want.

Upvotes: 0

Views: 619

Answers (2)

Ivan
Ivan

Reputation: 40768

Your mathematical formulas already gives you the outline to define the subscripts when using np.einsum. It is as simple as:

V = np.einsum('ij,ik->i', A, B)

which translates to V[i] = sum(A[i][k]*B[i][k]).

and

C = np.einsum('ik,ijk->ij', A, D)

i.e. C[i][j] = sum(A[i][k]*D[i][j][k]

You can read more about the einsum operator on this other question.

Upvotes: 4

hpaulj
hpaulj

Reputation: 231665

These can be done with element-wise multiplication and sum:

Vi = Σk Ai,k Bi,k
V = (A*B).sum(axis=-1)

and

Ci,j = Σk Ai,k Di,j,k
C = (A[:,None,:] * D).sum(axis=-1)

einsum may be faster, but it's a good idea of understand this approach.

To use matmul we have to add dimensions to fit the

(batch,i,j)@(batch,j,k) => (batch,k)

pattern. The sum-of-products dimension is j. Calculation-wise this is fast, but application here may be a bit tedious.

Upvotes: 1

Related Questions