Vectorized inner product with numpy

Question

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

Namely, in mathematical terms,

V_i = Σ_k A_i,k B_i,k

and

C_i,j = Σ_k A_i,k D_i,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.

thanks @juanpa.arrivillaga ! I hadn't found `einsum`, or, rather, I didn't click and read because `ein` was rather making me think of `eig`... If you make a very short answer twith this I'll accept it. — Joce, Jul 28 '21 at 10:29
Usually people who downvote don't stay around to see comments and respond. But my guess is that they'd like to see more detail, such as your "of course use a loop", and a [mcve]. — hpaulj, Jul 28 '21 at 16:19

Ivan · Accepted Answer · 2022-03-03T23:15:14.637

4

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.

edited Mar 03 '22 at 23:15

answered Jul 28 '21 at 10:40

Ivan

34,531
8
55
100

score 1 · Answer 2 · answered Jul 28 '21 at 15:49

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.

Vectorized inner product with numpy

2 Answers2