Understanding Numpy's einsum - Eli Bendersky's website

Why einsum?

numpy.einsum evaluates operations on multi-dimensional arrays using Einstein notation. Its explicit mode — where the subscript string includes -> and names the output dimensions — is common in ML papers. The notation is self-documenting: dimension labels such as i, j, or k make the intended shapes and contractions explicit in the code itself.

Matrix multiplication with explicit labels

Consider two matrices, A of shape (2,3) and B of shape (3,4):

>>> A = np.arange(6).reshape(2,3)

>>> A
array([[0, 1, 2],
       [3, 4, 5]])

>>> B = np.arange(12).reshape(3,4)+1

>>> B
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])

Multiplying them with einsum is done via the subscript string 'ij,jk->ik':

>>> np.einsum('ij,jk->ik', A, B)
array([[ 23,  26,  29,  32],
       [ 68,  80,  92, 104]])

The subscript string is a comma-separated list of input labels followed by -> and the output's label sequence. The letters are symbolic; actual dimension sizes are bound when the function is called with real arrays. Two rules govern the operation:

• The output labels determine the shape of the result.
• A label repeated among inputs but missing from the output is contracted — summed over.

Here, j appears in both inputs but not in the output, so each output element [ik] is the dot product of row i of the first matrix and column k of the second. Flipping the output labels to 'ij,jk->ki' yields the transpose:

>>> np.einsum('ij,jk->ki', A, B)
array([[ 23,  68],
       [ 26,  80],
       [ 29,  92],
       [ 32, 104]])

Batched multiplication and dimension reordering

The advantage of labeled notation becomes clear when working with higher-dimensional arrays. For a batch of matrices, einsum's explicit labels prevent ambiguity about which axes are being contracted. Given batched arrays Ab and Bb with a leading batch dimension:

>>> Ab = np.arange(6*6).reshape(6,2,3)
>>> Bb = np.arange(6*12).reshape(6,3,4)

The batched matmul Ab @ Bb “just works” in NumPy. With einsum, the same result is expressed as:

>>> np.einsum('bmd,bdn->bmn', Ab, Bb)

Here, b is the batch dimension and d is contracted. Note that b appears in the output, so it is not summed over. Output order is also under full control. The Fast Transformer Decoding paper (Noam Shazeer) defines tensors for batched multi-head attention — M of shape (b,m,d) and Pk of shape (h,d,k) — and computes per-head values in one step:

>>> m = 4; d = 3; k = 6; h = 5; b = 10
>>> Pk = np.random.randn(h, d, k)
>>> M = np.random.randn(b, m, d)

>>> np.einsum('bmd,hdk->bhmk', M, Pk).shape
(10, 5, 4, 6)

This subscript contracts d and orders the output as batch, head, sequence, and head size. Without explicit labels, the equivalent M @ Pk would be far less readable. The order of output axes is arbitrary; flipping to bhkm or another arrangement is trivial:

>>> np.einsum('bmd,hdk->hbmk', M, Pk).shape
(5, 10, 4, 6)

Multiple contractions and input transposes

einsum can contract more than one axis simultaneously. The same paper uses a subscript like:

>>> b = 10; n = 4; d = 3; v = 6; h = 5
>>> O = np.random.randn(b, h, n, v)
>>> Po = np.random.randn(h, d, v)
>>> np.einsum('bhnv,hdv->bnd', O, Po).shape
(10, 4, 3)

Here, h and v both appear in the two inputs but not in the output — each result element is a sum over both those axes. This is cumbersome to express otherwise. einsum can also transpose an input directly by writing its labels in the desired order. Multiplying A by its own transpose, which would be A @ A.T, is written as:

>>> np.einsum('ij,kj->ik', A, A)
array([[ 5, 14],
       [14, 50]])

The second operand is labeled kj rather than jk; the repeated label j is what gets contracted.

Chained operations

More than two operands are supported. Chaining a third matrix C into a product alongside A and B:

>>> C = np.arange(20).reshape(4, 5)

>>> A @ B @ C
array([[ 900, 1010, 1120, 1230, 1340],
       [2880, 3224, 3568, 3912, 4256]])

can be written in a single einsum call:

>>> np.einsum('ij,jk,kp->ip', A, B, C)
array([[ 900, 1010, 1120, 1230, 1340],
       [2880, 3224, 3568, 3912, 4256]])

A walkthrough implementation

A simplified mental model often suffices: the output labels define the loop structure, and any input label absent from the output becomes a summation axis. This can be turned directly into Python code. For the basic case 'ij,jk->ik', define dimension sizes and an empty output array:

i_size = __a.shape[0]
j_size = __a.shape[1]
assert j_size == __b.shape[0]
k_size = __b.shape[1]

out = np.zeros((i_size, k_size))

Looping over each output element and summing over the contracted axis j:

for i in range(i_size):
    for k in range(k_size):
        ...
return out

for i in range(i_size):
    for k in range(k_size):
        for j in range(j_size):
            out[i, k] += __a[i, j] * __b[j, k]
return out

The subscript fully determines the loop body, just as in Einstein notation. For multiple contractions, such as the subscript from the transformer example, the structure extends naturally:

for b in range(b_size):
    for n in range(n_size):
        for d in range(d_size):
            for v in range(v_size):
                for h in range(h_size):
                    out[b, n, d] += __a[b, h, n, v] * __b[h, d, v]

When there is no repeated label missing from the output, no summation loop is needed — the computation reduces to a product assignment. For 'i,j->ij', the result is the outer product of two 1D inputs:

def calc(__a, __b):
    i_size = __a.shape[0]
    j_size = __b.shape[0]

    out = np.zeros((i_size, j_size))

    for i in range(i_size):
        for j in range(j_size):
            out[i, j] = __a[i] * __b[j]
    return out

A full translation tool that emits the Python implementation text from an arbitrary subscript is available as translate_einsum in the associated GitHub repository.

Implicit mode

In implicit mode, the subscript omits -> and the output labels. Output shape is inferred by sorting input labels lexicographically. For standard matrix multiplication:

>>> np.einsum('ij,jk', A, B)
array([[ 23,  26,  29,  32],
       [ 68,  80,  92, 104]])

To obtain the transposed result, rearranging input labels suffices:

>>> np.einsum('ij,jh', A, B)
array([[ 23,  68],
       [ 26,  80],
       [ 29,  92],
       [ 32, 104]])

Because h precedes i lexicographically, this is equivalent to the explicit form 'ij,jh->hi'. Implicit mode is rarely seen in ML work — the legibility gained from explicit output labels is generally worth the small extra typing.

Historical context

Einstein introduced this notation in his 1916 paper on general relativity, where cumbersome nested summations over tensor components were common. A repeated index in a product implies summation over that index, regardless of whether it is written in subscript or superscript form. At the time, matrix notation was not yet popular in physics (Heisenberg introduced it in 1925), and Einstein's notation naturally extends to any number of dimensions, whereas matrix notation is primarily useful in 2D. The implicit mode of einsum is conceptually the closest to Einstein's original notation.

Dimension labels in a subscript can be any single letter, and input dimension sizes are checked for compatibility when a label is repeated across multiple operands.