带有 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.

Answer 1

Your mathematical formulas already give you the outline to define the subscript 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)

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

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

Answer 2

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.