What does np.einsum('mk,nk', D, D) do?

Question

I'm reading over someone else's code and am unsure what np.einsum does in this case.

print(np.einsum('mk,nk', D, D)) # D is an np array with shape (3, 100)

This code outputs an array with shape (3, 3).

Any help is appreciated!

Answer 1

There are four steps to convert your equation to einsum notation. Lets take this equation as an example C[i,k] = sum_j A[i,j] * B[j,k]

First we drop the variable names. We get ik = sum_j ij * jk We drop the sum_j term as it is implicit. We get ik = ij * jk We replace * with ,. We get ik = ij, jk The output is on the RHS and is separated with -> sign. We get ij, jk -> ik The einsum interpreter just runs these 4 steps in reverse. All indices missing in the result are summed over.

Here are some more examples from the docsstrong text

# Matrix multiplication
einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Trace
einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]

# Batch matrix multiplication
einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]