\( \newcommand{\vect}[1]{\boldsymbol{\mathbf{#1}}} \newcommand{\mat}[1]{\boldsymbol{\mathbf{#1}}} \newcommand{\R}{\mathbb{R}} \newcommand{\E}{\mathbb{E}} \newcommand{\Prob}{\mathbb{P}} \newcommand{\T}{^{\top}} \newcommand{\norm}[1]{\lVert #1 \rVert} \newcommand{\abs}[1]{\lvert #1 \rvert} \newcommand{\ihat}{\mkern1mu\hat{\boldsymbol{\imath}}\mkern2mu} \newcommand{\jhat}{\mkern1mu\hat{\boldsymbol{\jmath}}\mkern2mu} \newcommand{\pd}[2]{\frac{\partial #1}{\partial #2}} \newcommand{\argmax}{\operatorname*{arg\,max}} \newcommand{\argmin}{\operatorname*{arg\,min}} \newcommand{\cdata}[1]{\class{c-data}{#1}} \newcommand{\cmodel}[1]{\class{c-model}{#1}} \newcommand{\ctruth}[1]{\class{c-truth}{#1}} \newcommand{\cgrad}[1]{\class{c-grad}{#1}} \newcommand{\closs}[1]{\class{c-loss}{#1}} \newcommand{\cparam}[1]{\class{c-param}{#1}} \newcommand{\csecond}[1]{\class{c-second}{#1}} \renewcommand{\,}{\mkern3mu} \)

einsum: any contraction from scratch

Turn an index expression into a subscript string such as ij,jk->ik, build einsum from scratch with one loop over every index, then use numpy.einsum, batched matmul and einsum_path, and see why the order of a chain of products can change its cost a hundredfold.

M03.10.3L2 Intermediate25 min

In review (what this means)

The maths and code on this page are checked by automated tests on every change, but an independent reader has not signed it off yet. Spot a mistake or a confusing step? Tell us.

Before this lesson: M03.10.2 Index notation and the summation convention P02.2 Broadcasting and vectorization

What you will be able to do

• Translate between index expressions and einsum subscript strings, in explicit and in implicit mode.
• Implement a general contraction from scratch with itertools.product, and check it against numpy.einsum.
• Predict the output shape and the number of multiply-adds of a contraction from the sizes of its letters.
• Use the ellipsis, numpy.matmul and the @ operator for batched products, and einsum_path to choose the order of a chain.

The question

One NumPy function, np.einsum, computes a matrix product, a trace, a transpose, a dot product, a batch of matrix products and the attention scores of a transformer, each from a short string of letters. How can one loop do all of that, and why does the same product of three matrices take a hundred times more arithmetic when the string is evaluated in the wrong order?

From the last lesson

1. In \({c_{ik} = a_{ij} b_{jk}}\), which indices are free and which is the dummy?

Show answer

\(i\) and \(k\) are free: they appear once and name the entry of the result. \(j\) is the dummy: it appears twice and is summed.

2. What is the gradient of \(\vect{x}\T\mat{S}\vect{x}\), and when is it \(2\mat{S}\vect{x}\)?

Show answer

\({(\mat{S} + \mat{S}\T)\vect{x}}\). It equals \(2\mat{S}\vect{x}\) when \(\mat{S}\) is symmetric.

Intuition

Take an index expression, such as \({c_{ik} = a_{ij} b_{jk}}\), and keep only the letters: the subscripts of the inputs, separated by a comma, then an arrow, then the subscripts of the result. That gives the string ij,jk->ik. It holds everything the computation needs.

M03.10.3-F1

The anatomy of a subscript string. Letters kept in the output label the result; a letter that is in the inputs but not in the output is summed. Each letter keeps its colour on every axis it names, as in the explorer below.

The string is read by one rule. A letter in the output is kept; a letter that appears only in the inputs is summed. To compute it, loop over every combination of values of all the letters, multiply the matching entries of the inputs, and add the product into the output entry named by the output letters. One loop, any contraction.

Build it up

Step 1. Split the string. ij,jk->ik splits at the arrow into the inputs ij,jk and the output ik, and the inputs split at the comma into one term per operand.

Step 2. Read the size of each letter. Pair each term with its operand’s shape: ij with (2, 3) gives \(i \to 2\) and \(j \to 3\); jk with (3, 2) gives \(j \to 3\) and \(k \to 2\). A letter that appears twice must find the same size both times, or the string does not fit the arrays.

Step 3. Loop over every combination. The letters \(i\), \(j\), \(k\) take \({2 \cdot 3 \cdot 2} = 12\) combinations of values. itertools.product lists them.

Step 4. Multiply and add. For each combination, multiply \({A[i, j] \cdot B[j, k]}\) and add it into \({C[i, k]}\). Nothing else happens. A letter not in the output, here \(j\), changes from step to step while the output address stays put, so its terms pile up into one sum.

Step 5. A letter in every operand and in the output rides along. In bij,bjk->bik the letter \(b\) is in both inputs and in the output, so it is never summed. For each \(b\) the loop does a separate matrix product.

M03.10.3-F2

A batch axis rides along: one matrix product per value of \(b\). np.matmul and @ do the same for every leading axis, and broadcast those axes by the rule of P02.2.3.

The formal statement

NoteDefinition: the contraction of a subscript string

Let the inputs \(X^{(1)}, \dots, X^{(m)}\) carry the terms \(t_1, \dots, t_m\) and let the output term be \(o\). Every letter has one size, read from the shapes. Then

\[ \begin{aligned} \text{out}[o] = \sum X^{(1)}[t_1] \cdots X^{(m)}[t_m] , \end{aligned} \]

where the sum runs over every letter that is not in \(o\), and each bracket reads the entry at the current values of the letters. A letter repeated inside one term, as in ii, reads only entries whose two indices agree: the diagonal.

Three conventions of np.einsum complete the picture.

• Explicit mode writes the output after ->. Its letters may be in any order, which is how ij->ji transposes.
• Implicit mode leaves out the arrow. The output is then every letter that appears exactly once, in alphabetical order: np.einsum("ij,jk", A, B) is the product, and np.einsum("ii", A) is the trace.
• The ellipsis ... stands for any number of leading axes, broadcast as in NumPy: "...ij,...jk->...ik" is a matrix product over a stack of any shape, which is what np.matmul and @ compute.

The cost. The loop of the definition runs once per combination of letter values, so for letters of sizes \(n_1, \dots, n_r\) it takes \(n_1 n_2 \cdots n_r\) steps. With two inputs each step is one multiply-add. For ij,jk->ik with shapes \((a, b)\) and \((b, c)\) that is \(abc\), the familiar cost of a matrix product.

Why the order of a chain matters. For three matrices \(\mat{A}\) (\(a \times b\)), \(\mat{B}\) (\(b \times c\)) and \(\mat{C}\) (\(c \times d\)), the two orders cost

\[ \begin{aligned} (\mat{A}\mat{B})\mat{C} &: \; abc + acd , \\ \mat{A}(\mat{B}\mat{C}) &: \; bcd + abd . \end{aligned} \]

The first line is \(abc\) for \(\mat{A}\mat{B}\), an \(a \times c\) matrix, then \(acd\) to multiply it by \(\mat{C}\); the second is the same count for the other order. With \(a = c = 10\) and \(b = d = 1000\):

\[ \begin{aligned} abc + acd &= 100{,}000 + 100{,}000 \\ &= 2 \times 10^{5} , \\ bcd + abd &= 10^{7} + 10^{7} \\ &= 2 \times 10^{7} . \end{aligned} \]

A factor of 100. The answers are equal (matrix products are associative); only the work differs.

M03.10.3-F5

Same product, two orders. Multiplying the two matrices that share the long side first leaves a tiny 10 by 10 intermediate; the other order builds a million-entry matrix on the way.

np.einsum does not reorder by default (optimize=False), so a three-operand string runs as one naive loop over all four letters. np.einsum_path searches for a good order of pairwise contractions and reports its cost; optimize=True makes np.einsum use it.

Worked example

Run the loop by hand for ij,jk->ik on the running matrices \(\mat{A}\), rows \((4, 1)\) and \((2, 3)\), and \(\mat{S}\), rows \((2, 1)\) and \((1, 2)\). Every letter has size 2, so there are \({2 \cdot 2 \cdot 2} = 8\) steps. The loop runs over \(i\), then \(j\), then \(k\), with \(k\) changing fastest. The table counts from 1, as in maths; the code below counts from 0.

M03.10.3-F3

The eight steps of the loop. Each output entry receives two terms, one for each value of the summed letter \(j\), and the result is \(\mat{A}\mat{S}\), rows \((9, 6)\) and \((7, 8)\).

Step by step for the entry \(c_{11}\): at \((1, 1, 1)\) it receives \({a_{11} s_{11}} = 8\); at \((1, 2, 1)\) it receives \({a_{12} s_{21}} = 1\), making 9. The other three entries fill the same way, two terms each, and the result matches lesson M03.10.2.

A second string, ii->. Only the letter \(i\) appears, so the loop has 2 steps. The repeated letter reads \(a_{11} = 4\) and then \(a_{22} = 3\), and since \(i\) is not in the output both are added: the trace, 7.

Code

The definition, line for line: split, read sizes, loop, multiply, add.

from itertools import product
import numpy as np

def einsum_scratch(subscripts, *operands):
    """A general contraction: explicit mode with '->', or implicit mode without it."""
    subscripts = subscripts.replace(" ", "")
    if "->" in subscripts:
        inputs, output = subscripts.split("->")
    else:
        inputs = subscripts
        letters_in = inputs.replace(",", "")
        output = "".join(sorted(ch for ch in set(letters_in) if letters_in.count(ch) == 1))
    terms = inputs.split(",")
    size = {}
    for term, x in zip(terms, operands):
        for letter, n in zip(term, x.shape):
            if size.setdefault(letter, n) != n:
                raise ValueError(f"letter {letter} has sizes {size[letter]} and {n}")
    letters = list(dict.fromkeys(inputs.replace(",", "")))      # in order of first appearance
    out = np.zeros(tuple(size[ch] for ch in output))
    for values in product(*(range(size[ch]) for ch in letters)):
        at = dict(zip(letters, values))
        term_value = 1.0
        for term, x in zip(terms, operands):
            term_value *= x[tuple(at[ch] for ch in term)]
        out[tuple(at[ch] for ch in output)] += term_value
    return out

A = np.array([[4.0, 1.0], [2.0, 3.0]])
S = np.array([[2.0, 1.0], [1.0, 2.0]])
print("ij,jk->ik:", einsum_scratch("ij,jk->ik", A, S).tolist())
print("ii->:", einsum_scratch("ii->", A), "  ij->ji:", einsum_scratch("ij->ji", A).tolist())
print("i,i->:", einsum_scratch("i,i->", np.array([1.0, 2, 3]), np.array([4.0, 5, 6])))

ij,jk->ik: [[9.0, 6.0], [7.0, 8.0]]
ii->: 7.0   ij->ji: [[4.0, 2.0], [1.0, 3.0]]
i,i->: 32.0

Now against np.einsum on random arrays, including a batch, the implicit mode and the attention string of the next lesson.

rng = np.random.default_rng(0)
X, Y = rng.normal(size=(3, 2, 4)), rng.normal(size=(3, 4, 5))
Q, K = rng.normal(size=(2, 2, 3, 4)), rng.normal(size=(2, 2, 3, 4))
tests = [("bij,bjk->bik", X, Y), ("bhtd,bhsd->bhts", Q, K), ("ij,jk", A, S), ("ji", A), ("ij,jh", A, S)]
for s, *ops in tests:
    same = np.allclose(einsum_scratch(s, *ops), np.einsum(s, *ops))
    print(f"{s:16s} shape {np.einsum(s, *ops).shape}  agrees: {same}")
print("matmul and @ agree:", np.allclose(np.einsum("bij,bjk->bik", X, Y), X @ Y), np.matmul(X, Y).shape)
print("broadcast stacks:", (np.ones((4, 2, 3)) @ np.ones((3, 5))).shape, (np.ones((5, 1, 2, 3)) @ np.ones((4, 3, 6))).shape)
print("ij,jh is the transpose of AS:", np.einsum("ij,jh", A, S).tolist())

bij,bjk->bik     shape (3, 2, 5)  agrees: True
bhtd,bhsd->bhts  shape (2, 2, 3, 3)  agrees: True
ij,jk            shape (2, 2)  agrees: True
ji               shape (2, 2)  agrees: True
ij,jh            shape (2, 2)  agrees: True
matmul and @ agree: True (3, 2, 5)
broadcast stacks: (4, 2, 5) (5, 4, 2, 6)
ij,jh is the transpose of AS: [[9.0, 7.0], [6.0, 8.0]]

Finally the chain of three matrices. np.einsum_path reports a FLOP count that counts each multiplication and each addition, so it is twice the number of multiply-adds for a pairwise step.

import re

A3, B3, C3 = rng.normal(size=(10, 1000)), rng.normal(size=(1000, 10)), rng.normal(size=(10, 1000))
path, info = np.einsum_path("ij,jk,kl->il", A3, B3, C3, optimize="greedy")
_, other = np.einsum_path("ij,jk,kl->il", A3, B3, C3, optimize=["einsum_path", (1, 2), (0, 1)])
flops = lambda text: float(re.search(r"Optimized FLOP count:\s*([0-9.e+]+)", text).group(1))
print("chosen path:", path, "  FLOPs:", flops(info), "  the other order:", flops(other))
print("ratio:", flops(other) / flops(info))
print("same answer:", np.allclose(np.einsum("ij,jk,kl->il", A3, B3, C3, optimize=True), A3 @ B3 @ C3))

chosen path: ['einsum_path', (0, 1), (0, 1)]   FLOPs: 400000.0   the other order: 40000000.0
ratio: 100.0
same answer: True

The path (0, 1) contracts operands 0 and 1 first, \(\mat{A}\mat{B}\), as the formal section predicted. PyTorch’s torch.einsum takes the same subscript strings; the course does not run it, because the deep-learning library comes in DL01.7.

Explore

M03.10-X1

einsum explorer

Press the presets one by one and read which letters the explorer marks as summed. Then change the sizes: which sizes change the output shape, and which change only the cost? Finally type your own strings, including broken ones such as ij,jk->il or ij,jk->ii, and read the reason.

Pitfalls

M03.10.3-F4

Wrong beside right. A letter that appears in the output is never summed, so ij,ij->ij is the entry-by-entry product, not the matrix product.

“A letter shared by two inputs is always summed.”

Only if it is missing from the output. ij,ij->ij keeps both letters, so it multiplies entry by entry; ij,ij-> sums every product (17 for \(\mat{A}\) and \(\mat{S}\)); ij,jk->ik sums over \(j\) alone. The output decides.

“In implicit mode the output keeps the order I wrote.”

It is sorted alphabetically. np.einsum("ij,jh", A, S) has output letters \(h\) and \(i\), sorted as hi, so it returns the transpose of \(\mat{A}\mat{S}\): rows \((9, 7)\) and \((6, 8)\). Write the arrow and the output whenever the order matters, which is almost always.

“einsum always finds the fastest order.”

np.einsum has optimize=False by default and evaluates a string with three or more operands in one naive loop: for the chain above that is \({10 \cdot 1000 \cdot 10 \cdot 1000} = 10^8\) steps. Pass optimize=True, or check the plan with np.einsum_path.

Check yourself

Recall

1. Write the subscript string for \({y_i = a_{ij} x_j}\).

Show answer

ij,j->i.

2. Which letters of a subscript string are summed?

Show answer

Those that appear in the inputs but not in the output.

3. What output does np.einsum("ij,jk", A, B) produce in implicit mode, and why?

Show answer

\(\mat{A}\mat{B}\). Without an arrow the output is every letter that appears exactly once, \(i\) and \(k\), in alphabetical order, and \(j\) appears twice, so it is summed.

4. How many steps does the loop take for bij,bjk->bik with sizes \(b = 3\), \(i = 2\), \(j = 4\), \(k = 5\)?

Show answer

\({3 \cdot 2 \cdot 4 \cdot 5} = 120\), the product of the sizes of all the letters.

5. What do np.matmul and @ do with axes before the last two?

Show answer

They treat them as a stack of matrices and broadcast them, multiplying the last two axes as matrices: (4, 2, 3) @ (3, 5) has shape (4, 2, 5).

Exercises

1. ★ Write the subscript string for each: the outer product of two vectors, the sum of all entries of a matrix, the row sums of a matrix, and the matrix–vector product with \(\mat{A}\T\).

Show solution

Outer product i,j->ij (nothing summed). Sum of all entries ij-> (both letters summed). Row sums ij->i (only \(j\) summed). \(\mat{A}\T\vect{x}\) is \({y_j = a_{ij} x_i}\), so ij,i->j.

2. ★★ A batched product bij,bjk->bik runs with \(b = 32\), \(i = 64\), \(j = 128\), \(k = 64\). Give the output shape and the number of multiply-adds, and say which letter you could double without changing the output shape.

Show solution

The output keeps \(b\), \(i\), \(k\): shape (32, 64, 64). The loop runs over all four letters: \({32 \cdot 64 \cdot 128 \cdot 64} = 16{,}777{,}216\) multiply-adds. Doubling \(j\), the summed letter, doubles the cost and leaves the shape unchanged.

3. ★★★ Let \(\mat{A}\) and \(\mat{B}\) be \(n \times n\) and \(\vect{x}\) a vector of length \(n\). Count the multiply-adds of \((\mat{A}\mat{B})\vect{x}\) and of \(\mat{A}(\mat{B}\vect{x})\), and find their ratio for \(n = 1000\).

Show solution

\(\mat{A}\mat{B}\) is ij,jk->ik with three letters of size \(n\): \(n^3\) steps. Then \((\mat{A}\mat{B})\vect{x}\) is ik,k->i: \(n^2\). Total \(n^3 + n^2\).

\(\mat{B}\vect{x}\) is jk,k->j: \(n^2\), and \(\mat{A}(\mat{B}\vect{x})\) another \(n^2\). Total \(2n^2\).

For \(n = 1000\) the ratio is

\[ \begin{aligned} \frac{n^3 + n^2}{2n^2} = \frac{n + 1}{2} = 500.5 . \end{aligned} \]

Multiply by the vector first. np.einsum_path("ij,jk,k->i", ...) finds this order.

Summary card

Kept letters label the result; missing letters are summed; one loop does it all.

M03.10.3-F1

ij,jk->ik is \(\mat{A}\mat{B}\).

• Rule: multiply the matching entries; sum over every letter missing from the output.
• Implicit mode: output = letters appearing once, sorted alphabetically.
• Cost: product of the sizes of all letters; two inputs, one multiply-add per step.
• Chains: \((\mat{A}\mat{B})\mat{C}\) costs \(abc + acd\), \(\mat{A}(\mat{B}\mat{C})\) costs \(bcd + abd\); use einsum_path or optimize=True.

Sources and next steps

Sources. Each was opened and checked against this lesson on 15 November 2026.

• NumPy reference, numpy.einsum: implicit mode reorders the output axes alphabetically (so np.einsum('ji', a) transposes and np.einsum('ij,jh', a, b) returns the transpose of the product), 'ii' is the trace and 'ii->i' the diagonal, the ellipsis enables broadcasting, and optimize defaults to False (True means the greedy algorithm).
• NumPy reference, numpy.einsum_path: the 'greedy' and 'optimal' searches, the path as a list of pairs of operand positions, and the printed naive and optimised FLOP counts.
• NumPy reference, numpy.matmul: an argument with more than two axes is treated as a stack of matrices in the last two indexes and broadcast accordingly; matmul implements the @ operator.
• PyTorch reference, torch.einsum: the same subscript strings (letters for dimensions, commas between operands, -> for the output, alphabetical order in implicit mode), with bij,bjk->bik as its batch matrix product.

Next. M03.10.4 Tensors inside a network layer reads a dense layer, multi-head attention scores and a low-rank weight matrix as contractions.

This lesson unlocks

M03.10.4 Tensors inside a network layer

Your progress is saved in this browser only. No account, no tracking.