NumPy Einsum — Core Concepts

Why this topic matters

np.einsum (Einstein summation) is one of the most versatile functions in NumPy. It can express dot products, matrix multiplication, traces, transposes, outer products, batch operations, and tensor contractions — all with a single subscript string. Learning it replaces a grab bag of specialized functions with one consistent notation.

How the subscript string works

The general form is 'input_subscripts->output_subscripts':

• Each operand gets comma-separated subscript labels (one letter per dimension).
• Repeated labels between operands imply element-wise pairing along that axis.
• Labels present in the input but missing from the output are summed over.
• Labels in the output define the shape of the result.

Common operations as einsum

Operation	NumPy call	Einsum
Vector dot product	np.dot(a, b)	'i,i->'
Matrix multiplication	a @ b	'ij,jk->ik'
Element-wise multiply	a * b	'ij,ij->ij'
Transpose	a.T	'ij->ji'
Trace	np.trace(a)	'ii->'
Diagonal	np.diag(a)	'ii->i'
Outer product	np.outer(a, b)	'i,j->ij'
Sum all	a.sum()	'ij->'
Sum over rows	a.sum(axis=0)	'ij->j'
Batch matrix multiply	np.matmul(A, B)	'bij,bjk->bik'

Reading subscripts step by step

Take 'ij,jk->ik' (matrix multiplication):

1. First operand has dimensions i (rows) and j (cols).
2. Second operand has dimensions j (rows) and k (cols).
3. j appears in both inputs but not in the output → sum over j.
4. Output has dimensions i and k → result is (i × k).

This is exactly the definition of matrix multiplication: for each (i, k), sum over j of A[i,j] * B[j,k].

Implicit vs explicit mode

If you omit the ->, einsum uses implicit mode: output labels are the sorted unique labels not summed over:

# Implicit — einsum decides output subscripts
np.einsum('ij,jk', A, B)  # same as 'ij,jk->ik'

# Explicit — you specify output
np.einsum('ij,jk->ik', A, B)  # same result

Explicit mode is always clearer. Use it.

Common misconception

People assume einsum is always faster than specialized NumPy functions. It is not. For standard operations (matmul, dot), NumPy dispatches to optimized BLAS libraries. Einsum may or may not find the same path. Where einsum wins is complex multi-step operations that would otherwise require temporary arrays:

# Without einsum: two temporaries
temp = A * B       # (n, n) temporary
result = temp.sum(axis=1)  # another pass

# With einsum: no temporaries
result = np.einsum('ij,ij->i', A, B)  # one pass, less memory

The optimize parameter

For expressions with three or more operands, the order of contraction matters enormously. Pass optimize=True (or optimize='optimal') to let NumPy find the best contraction order:

# Without optimize: might create huge intermediates
result = np.einsum('ijk,jkl,lm->im', A, B, C, optimize=True)

For two-operand expressions, optimization has no effect — there is only one possible order.

The one thing to remember: Einsum’s subscript string is a mini language for array math — learn to read it like a formula and you will never fumble with axis arguments again.