Topolearneinsum makes complex tensor contractions readable and avoids building huge intermediate arrays.
np.einsum(subscripts, *operands) describes an operation by labelling each axis with a letter. The rules:
- Letters before
->label the axes of each input, separated by commas. - Letters after
->are the output axes, in that order. - Any letter that appears in the inputs but not in the output is summed over.
- A letter repeated across inputs means those axes are multiplied element-wise (they must have equal length).
So "ij,jk->ik" is matrix multiplication: j is shared and dropped, hence multiplied and summed. A handful of translations make the idea click:
"ij->ji"transpose;"ii->i"diagonal;"ii->"trace"i,i->"dot product;"i,j->ij"outer product"ij,ij->i"row-wise dot products, without forming the full product matrix"bij,bjk->bik"batched matrix multiply"...ij,...jk->...ik"the same with any number of leading batch axes
Why use it? Readability for tensor code (attention mechanisms, physics, statistics), and memory: einsum("ij,ij->i", A, B) computes row dots without allocating A * B.
With three or more operands, the order of pairwise contractions can change cost by orders of magnitude. Pass optimize=True (or "optimal") to let NumPy pick a good order and dispatch to BLAS where possible; np.einsum_path shows the plan and estimated cost. For plain two-operand matrix products, @ is usually at least as fast because it calls BLAS directly; einsum's own loop may not.
Example
import numpy as np
A = np.arange(6.0).reshape(2, 3)
B = np.arange(12.0).reshape(3, 4)
print(np.allclose(np.einsum("ij,jk->ik", A, B), A @ B)) # True
M = np.arange(9.0).reshape(3, 3)
print(np.einsum("ii->", M), np.trace(M)) # 12.0 12.0
print(np.einsum("ii->i", M)) # [0. 4. 8.]
X = np.arange(6.0).reshape(3, 2)
Y = np.ones((3, 2))
print(np.einsum("ij,ij->i", X, Y)) # row dots [1. 5. 9.]
batch = np.ones((10, 2, 3)); w = np.ones((10, 3, 4))
print(np.einsum("bij,bjk->bik", batch, w).shape) # (10, 2, 4)
# three operands: let NumPy choose the contraction order
a, b, c = np.ones((50, 400)), np.ones((400, 400)), np.ones((400, 5))
r = np.einsum("ij,jk,kl->il", a, b, c, optimize=True)
print(r.shape) # (50, 5)
path, info = np.einsum_path("ij,jk,kl->il", a, b, c, optimize="optimal")
print(path) # contraction order chosenRemember
- Each letter labels an axis; letters missing from the output are summed
- Shared letters across inputs multiply element-wise
...stands for any leading batch axes- Use
optimize=Truewith 3+ operands;@is often faster for one matmul
Watch out for
- Forgetting
->and relying on implicit mode, where output axes are the alphabetically sorted unrepeated letters, which can transpose your result unexpectedly.