The Joy of einsum

NumPy's np.einsum uses Einstein summation notation to express tensor operations in a single, readable string. Once you learn the notation, you'll never want to go back to juggling transpose, reshape, tensordot, np.sum(axis=...), and np.diagonal again.

1. Matrix Multiply

The Old Way

# option A
C = np.dot(A, B)

# option B
C = A @ B

# option C
C = np.tensordot(A, B, axes=1)

# option D — for the truly unhinged
C = np.sum(
    A[:, :, np.newaxis] * B[np.newaxis, :, :],
    axis=1
)

einsum

C = np.einsum('ik,kj->ij', A, B)

# Read it aloud:
#   "for each i,j: sum over k
#    of A[i,k] * B[k,j]"
#
# The subscript string IS the math.

2. Batch Matrix Multiply

The Old Way

# np.matmul handles batches...
C = np.matmul(A, B)      # or A @ B

# but tensordot? you need a loop:
C = np.stack([
    np.tensordot(A[b], B[b], axes=1)
    for b in range(A.shape[0])
])

einsum

C = np.einsum('bik,bkj->bij', A, B)

# Just add 'b'. Done.

3. Outer Product

The Old Way

M = np.outer(a, b)

# or the broadcasting trick
M = a[:, np.newaxis] * b[np.newaxis, :]

einsum

M = np.einsum('i,j->ij', a, b)

# No repeated index = no sum.
# Just pair every i with every j.

4. Trace

The Old Way

t = np.trace(A)

# or
t = np.sum(np.diagonal(A))

einsum

t = np.einsum('ii->', A)

# Repeated index, no output index
# = sum the diagonal. Two characters.

5. Bilinear Form

The Old Way

result = x @ M @ y

# sure, that works for 1-D vectors.
# but for a batch of vectors?
result = np.sum(
    x[:, :, np.newaxis]
    * (M @ y[:, :, np.newaxis]),
    axis=1
).squeeze()

einsum

# single
result = np.einsum('i,ij,j->', x, M, y)

# batched
result = np.einsum('bi,ij,bj->b', x, M, y)

# The formula IS the code.

6. Attention Scores

The Old Way

# Q, K shape: (batch, heads, seq, dim)

scores = np.matmul(
    Q,
    np.transpose(K, (0, 1, 3, 2))
)

# wait, was it (0,1,3,2) or (0,2,1,3)?
# *checks docs for the 50th time*

einsum

# Q, K shape: (batch, heads, seq, dim)

scores = np.einsum(
    'bhsd,bhtd->bhst', Q, K
)

# b=batch, h=heads, s/t=seq, d=dim
# contract over d, keep everything else.
# no transpose axes to memorize.

7. Tensor Contraction

Contract a 4-D tensor with a 3-D tensor over two shared axes. Good luck without einsum.

The Old Way

# T: (a,b,c,d)  U: (c,d,e)
# contract over c and d

R = np.tensordot(T, U, axes=[(2,3), (0,1)])

# Quick: which axes map to which?
# Is the output (a,b,e) or (e,a,b)?
# Hope you like reading docs!

einsum

# T: (a,b,c,d)  U: (c,d,e)

R = np.einsum('abcd,cde->abe', T, U)

# output shape? it's right there: abe
# contracted axes? c and d (repeated)
# nothing to memorize.

The Rosetta Stone

Operation	einsum	Alternative
Vector dot product	einsum('i,i->', a, b)	np.dot(a, b)
Outer product	einsum('i,j->ij', a, b)	np.outer(a, b)
Element-wise product	einsum('ij,ij->ij', A, B)	A * B
Matrix-vector product	einsum('ij,j->i', A, x)	A @ x
Matrix multiply	einsum('ik,kj->ij', A, B)	A @ B
Transpose	einsum('ij->ji', A)	A.T
Trace	einsum('ii->', A)	np.trace(A)
Sum all	einsum('ij->', A)	np.sum(A)
Sum over rows	einsum('ij->j', A)	np.sum(A, axis=0)
Sum over cols	einsum('ij->i', A)	np.sum(A, axis=1)
Diagonal	einsum('ii->i', A)	np.diagonal(A)
Batch matmul	einsum('bik,bkj->bij', A, B)	np.matmul(A, B)
Bilinear form	einsum('i,ij,j->', x, M, y)	x @ M @ y
Quadratic form	einsum('i,ij,j->', x, M, x)	x @ M @ x
Frobenius norm²	einsum('ij,ij->', A, A)	np.sum(A**2)
Batch outer product	einsum('bi,bj->bij', a, b)	a[:,:,None]*b[:,None,:]

The Three Rules of einsum

1. Repeated indices are summed over. An index that appears in the inputs but not in the output is contracted (summed out). 'ik,kj->ij' — k is repeated and absent from the output, so we sum over it.

2. Output indices determine the shape. The indices after -> specify the axes (and their order) in the result. 'ij->ji' is a transpose; 'ij->i' sums out j.

3. Free indices are broadcast. If an index appears in one input but not the other, it's just carried along. 'bi,ij->bj' — b only appears in the first input, acting as a batch dimension.

Pro Tips

Use optimize=True for big contractions

np.einsum('abc,bcd,cde->ae', X, Y, Z, optimize=True)
NumPy will find an efficient contraction order automatically (via np.einsum_path). For chains of 3+ tensors, this can be orders of magnitude faster.

It works everywhere

The same string notation works in np.einsum, torch.einsum, jnp.einsum (JAX), tf.einsum, and the standalone opt_einsum library. Learn it once, use it everywhere.

Name your indices meaningfully

Nothing stops you from using 'batch heads seq dim, batch heads dim seq -> batch heads seq seq' — but single letters like b h s d are conventional. Pick letters that remind you of the dimension: batch, heads, sequence, dim.

When NOT to use einsum

For simple A @ B matrix multiplies, the @ operator is more readable. Einsum shines when you'd otherwise need transpose + reshape + sum combos, or when working with 3+ dimensions.