Introduction

numpy.einsum is a compact way to describe tensor operations using index labels instead of a long chain of reshapes and matrix helpers. It looks unfamiliar at first because the syntax comes from Einstein summation notation, but the underlying rule is simple: name axes, decide which ones are multiplied together, and decide which ones remain in the output. Once that clicks, many array operations become easier to express.

Read the Subscript String as Axis Names

Consider this example:

1import numpy as np
2
3A = np.array([[1, 2], [3, 4]])
4x = np.array([10, 20])
5result = np.einsum("ij,j->i", A, x)
6print(result)

Read "ij,j->i" step by step:

• the first input has axes i and j
• the second input has axis j
• repeated labels are multiplied and summed over
• only i remains in the output

That is matrix-vector multiplication. einsum is not magic here. It is just describing the same operation in axis language.

Learn It Through Familiar Operations

Many common NumPy operations have direct einsum versions.

Dot product:

1import numpy as np
2
3a = np.array([1, 2, 3])
4b = np.array([4, 5, 6])
5print(np.einsum("i,i->", a, b))

Transpose:

1import numpy as np
2
3M = np.array([[1, 2, 3], [4, 5, 6]])
4print(np.einsum("ij->ji", M))

Matrix multiplication:

1import numpy as np
2
3A = np.array([[1, 2], [3, 4]])
4B = np.array([[5, 6], [7, 8]])
5print(np.einsum("ik,kj->ij", A, B))

These examples are worth memorizing because they build intuition. Once you can translate these, more advanced contractions feel much less abstract.

Repeated Labels Mean Reduction

The fastest way to understand einsum is to track which labels disappear. Any label that is repeated across inputs and not included in the output is summed away.

For example, row and column sums can be written directly:

1import numpy as np
2
3M = np.array([[1, 2], [3, 4]])
4print(np.einsum("ij->i", M))
5print(np.einsum("ij->j", M))
6print(np.einsum("ij->", M))

The first keeps rows, the second keeps columns, and the third keeps no axes at all, so it returns the total sum. Thinking in terms of labels that survive is often easier than thinking in terms of a memorized formula.

Batched Tensor Operations Are Where einsum Shines

einsum becomes especially valuable when the operation involves batch dimensions or more than two axes. Instead of reshaping by hand, you can write the contraction directly.

1import numpy as np
2
3X = np.random.rand(2, 3, 4)
4Y = np.random.rand(2, 4, 5)
5Z = np.einsum("bij,bjk->bik", X, Y)
6print(Z.shape)

Here b is the batch axis, so each item in the batch is multiplied independently. This is a good example of why einsum is popular in scientific computing and machine learning code: it exposes the intended axis behavior directly.

Use optimize=True for Larger Contractions

For multi-operand expressions, contraction order can matter. NumPy can choose a better path when you ask it to optimize.

1import numpy as np
2
3a = np.random.rand(10, 20)
4b = np.random.rand(20, 30)
5c = np.random.rand(30, 40)
6
7result = np.einsum("ab,bc,cd->ad", a, b, c, optimize=True)
8print(result.shape)

This does not mean einsum is always the fastest tool, but it does mean you should not judge performance from the syntax alone. For simple operations, matmul or sum may still be clearer.

Use einsum When It Clarifies the Operation

The best use of einsum is not showing off compactness. The best use is making a tensor relationship explicit. If a reader already understands A @ B, use @. If a contraction involves several axes and awkward reshapes, einsum may be the clearest representation.

A good engineering rule is to choose the most direct expression of the operation. einsum is powerful, but clarity still matters more than cleverness.

Common Pitfalls

• Treating the subscript string as magic instead of mapping each label to a real axis.
• Forgetting that repeated labels are usually summed away unless they appear in the output.
• Using einsum for very simple operations where @ or sum would be clearer.
• Mixing up label order and producing the right numbers in the wrong shape.
• Ignoring optimize=True when contracting several operands and then blaming einsum for poor performance.

Summary

• 'einsum names axes explicitly and describes which ones survive into the output.'
• Repeated labels typically indicate multiplication followed by summation.
• Familiar operations such as dot product, transpose, and matrix multiplication can all be written with einsum.
• The function is especially helpful for batched and higher-rank tensor contractions.
• Use it when it makes the operation clearer, not just shorter.