PyTorch Einsum

torch.einsum is a function that names tensor dimensions with letters inside a string. The letters say which dimensions multiply and which survive into the result.

Matmul in one string

The string ij,jk->ik shares the letter j between the two tables. That letter is multiplied and summed away, exactly like the @ operator.

import torch

A = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])
B = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0]])
by_einsum = torch.einsum("ij,jk->ik", A, B)
by_at = A @ B
print(by_einsum)
print("Equal:", torch.equal(by_einsum, by_at))

A batch letter rides along

The letter b appears in both operands and in the output. It is carried through, so each batch item multiplies its own pair of tables.

torch.manual_seed(0)
P = torch.randn(2, 3, 4)
R = torch.randn(2, 4, 5)
batched = torch.einsum("bij,bjk->bik", P, R)
print("Batch result shape:", batched.shape)

An outer product

No letter is shared in this string. Every value of a pairs with every value of b, and nothing is summed away.

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0])
outer = torch.einsum("i,j->ij", a, b)
print(outer)

A size mismatch fails

The shared letter j has size 3 in one operand and size 4 in the other. The call raises a RuntimeError, and its last line names the clashing sizes.

bad = torch.einsum("ij,jk->ik", torch.randn(2, 3), torch.randn(4, 2))

Example

import torch

Q = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0]])
K = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0],
                  [2.0, 0.0]])

scores = torch.einsum("qd,kd->qk", Q, K)

loop = torch.zeros(3, 4)
for i in range(3):
    for j in range(4):
        loop[i, j] = (Q[i] * K[j]).sum()

print("Einsum scores:")
print(scores)
print("Loop scores:")
print(loop)
print(f"Max difference: {(scores - loop).abs().max().item()}")

Key points

Subscript string: the letters in quotes that name one dimension of each operand
A letter shared by two operands is multiplied along that dimension and added up
A letter missing from the output side is summed away
The arrow -> separates the operand letters from the output letters
One notation covers matmul, batch matmul, and outer products
In short, einsum writes the dimension plan once and lets torch run the loops

In this example, Q and K are two small tables of vectors. The rows of Q are queries and the rows of K are keys, the two sides that attention compares. The string qd,kd->qk shares the letter d, the feature dimension. einsum multiplies every query with every key along d and adds the products up. The double loop below repeats the same maths one pair at a time. Both score tables print the same numbers, and the max difference is 0.0. The next lesson scales these scores, turns them into weights, and mixes value vectors with them.

Note

A mismatched letter size raises a RuntimeError before any maths runs

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

Q = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0]])
K = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0],
                  [2.0, 0.0]])

scores = torch.einsum("qd,kd->qk", Q, K)

loop = torch.zeros(3, 4)
for i in range(3):
    for j in range(4):
        loop[i, j] = (Q[i] * K[j]).sum()

print("Einsum scores:")
print(scores)
print("Loop scores:")
print(loop)
print(f"Max difference: {(scores - loop).abs().max().item()}")

Exercise

In the string 'ij,jk->ik', what happens to the letter j?

It is multiplied and summed away It survives into the output It names the batch dimension It is ignored