einsum and tensordot

The formula

You name each axis with a letter and specify what you want in the result. The rule is short: repeated letters between the inputs are multiplied , and letters thatdo not appear in the output are summed . Matrix multiplication is the classic example:multiplies by the "ij,jk->ik" (repeated) and sums it (because the j is not in the output): j With the same rule, the formula expresses dozens of operations:

import numpy as np

A = np.arange(6).reshape(2, 3)
B = np.arange(12).reshape(3, 4)
print(np.einsum("ij,jk->ik", A, B))
print(np.array_equal(np.einsum("ij,jk->ik", A, B), A @ B))

[[20 23 26 29]
 [56 68 80 92]]
True

A dictionary of patterns

Formula

M = np.arange(9).reshape(3, 3)
print(np.einsum("ii->", M), np.einsum("ii->i", M), np.einsum("ij->i", M), np.einsum("ij->j", M))

v = np.array([1, 2, 3])
w = np.array([4, 5, 6])
print(np.einsum("i,i->", v, w), np.einsum("i,j->ij", v, w))

12 [0 4 8] [ 3 12 21] [ 9 12 15]
32 [[ 4  5  6]
 [ 8 10 12]
 [12 15 18]]

The real use case: attention in neural networks

The central calculation of language models compares each "query" with each "key", in a batch. In formula, "bqd,bkd->bqk": for each example b, it multiplies the vector of dimension d of each query q by each key k. Without the einsum, you would need to transpose the last dimension manually:

rng = np.random.default_rng(0)
Q = rng.normal(size=(2, 5, 8))
K = rng.normal(size=(2, 7, 8))
scores = np.einsum("bqd,bkd->bqk", Q, K)
print(scores.shape, np.allclose(scores, Q @ K.transpose(0, 2, 1)))

(2, 5, 7) True

The formula documents the axes: each letter has a meaning, and the reader understands the operation without deciphering a sequence of transpose.

`tensordot` and the order of multiplications

The np.tensordot sum over the axes you indicate (the parameter axes), and for two matrices with axes=1 it is equivalent to @. The same operation can be much cheaper depending on the order in which a chain of matrices is multiplied, because multiplication is associative, but the cost is not:

import timeit

print(np.allclose(np.tensordot(A, B, axes=1), A @ B))

C1 = rng.normal(size=(2000, 3))
C2 = rng.normal(size=(3, 2000))
C3 = rng.normal(size=(2000, 3))
esquerda = lambda: (C1 @ C2) @ C3
direita = lambda: C1 @ (C2 @ C3)
print(np.allclose(esquerda(), direita()))
t_esq = min(timeit.repeat(esquerda, number=1, repeat=3))
t_dir = min(timeit.repeat(direita, number=1, repeat=3))
print("a ordem importa:", t_esq > 3 * t_dir)

True
True
a ordem importa: True

The first way creates a 2000 by 2000 matrix in the middle of the process, and the second reduces everything to a 3 by 3 matrix first. The result is the same, and the cost, very different. The np.linalg.multi_dot chooses the best order for you, and the einsum with optimize=True does the same.

When I don't use einsum

For a simple matrix multiplication, the @ is clearer and uses the optimized libraries directly. The einsum pays off when there are more than two inputs, batch axes, or unusual sums, and when the formula documents better than the equivalent code.