NumPy Einsum for Tensor Operations

Expressing complex multi-dimensional array operations clearly and efficiently is a common challenge in scientific computing and data science. NumPy's np.einsum function provides a powerful, declarative solution using a concise notation system. Mastering einsum allows you to replace verbose, error-prone loops with a single, readable expression for operations ranging from simple matrix multiplication to high-order tensor contractions, which are foundational in fields like deep learning and physics.

Understanding Einstein Summation Notation

The core of np.einsum is Einstein summation convention, a notational shorthand for summing over repeated indices. In NumPy's implementation, you define an operation using a subscript string and the input arrays. The function then performs the specified summations and products.

The general syntax is np.einsum(subscript_string, operand1, operand2, ...). The subscript string is composed of labels for each dimension of every input array, separated by commas, followed by an arrow -> and the labels for the output array's dimensions. The rule is simple: any index label that appears in the inputs but not in the output is summed over (or "contracted").

For example, consider the expression for the inner product of two vectors, $\vec{a}$ and $\vec{b}$. The mathematical operation is ${\sum }_i{a}_i{b}_i$. Notice the index $i$ is repeated and summed over. In einsum, this is written as 'i,i->'. The i label appears for both inputs, but is absent in the output (a scalar), so it is contracted. The code is result = np.einsum('i,i->', a, b).

Core Operations with Einsum

The true power of einsum becomes apparent when performing common linear algebra and tensor operations without reshaping or transposing arrays manually.

Matrix Multiplication: For two matrices $A$ and $B$, standard multiplication is $C_{ij}={\sum }_k{A}_{ik}{B}_{kj}$. The repeated index $k$ is summed. The corresponding einsum expression is 'ik,kj->ij'. This clearly shows which axes are contracted (k) and which form the output (i and j).

Trace: The trace of a matrix $A$ is the sum of its diagonal elements: ${\sum }_i{A}_{ii}$. Here, the index i is repeated in a single input, indicating a diagonal selection, and is absent from the output, meaning it's summed. The expression is simply 'ii->'.

Transpose: While A.T is simpler, einsum can reorder axes arbitrarily. To transpose a matrix, you just reorder the indices in the output: 'ij->ji'. For higher-dimensional arrays, like converting a tensor from shape (i,j,k) to (k,i,j), you would write 'ijk->kij'.

Batch Operations: A major strength is handling batch dimensions effortlessly. For batch matrix multiplication, you have a stack of matrices A with shape (b, i, k) and B with shape (b, k, j). You want to multiply each pair, resulting in shape (b, i, j). The einsum expression 'bik,bkj->bij' keeps the batch dimension b free (it appears on both sides of the arrow) while contracting over k$. This is far clearer than a manual loop or using np.matmul`.

Tensor Contractions and Deep Learning Applications

Tensor contractions generalize matrix multiplication to higher dimensions. They involve summing over the product of elements across shared indices of two or more tensors. For instance, contracting a 3D tensor $X_{ijk}$ with a 2D tensor $Y_{jl}$ over the index $j$ would be written mathematically as $Z_{ikl}={\sum }_j{X}_{ijk}{Y}_{jl}$. The einsum expression mirrors this perfectly: 'ijk,jl->ikl'.

This capability is directly applicable in deep learning gradient computations. Consider a common operation: calculating the gradient of a loss with respect to the weights in a fully connected layer. If you have an error tensor $\delta$ with shape (batch, outputnodes) and an input activation tensor MATHINLINE15 with shape (batch, inputnodes), the weight gradient is $\Delta W=A^T\cdot \delta$. Using einsum, this is expressed as 'bi,bo->io'$, where i are input nodes, o are output nodes, and b` is the batch dimension that gets summed over. This single line precisely computes the necessary outer products and sums across the batch, which is exactly what backpropagation requires.

Optimization and Performance Considerations

While einsum is expressive, it's important to understand its performance profile. For many operations, especially those with multiple contractions, NumPy's internal einsum logic may not choose the most computationally efficient order of operations.

This is where the optimize=True argument becomes crucial. When you call np.einsum('...', a, b, c, optimize=True), NumPy analyzes the subscript string to find the optimal contraction path—the order in which to pairwise contract tensors to minimize intermediate array size and total floating-point operations. For contractions involving three or more arrays, this can lead to orders-of-magnitude speed improvements.

How does einsum compare to other methods?

• Explicit Loops: einsum is implemented in C, making it vastly faster than Python loops for any non-trivial array size. It should always be preferred over manual iteration.
• np.matmul/@ operator: For standard 2D matrix multiplication or batch 2D multiplication, np.matmul is often slightly faster and is the most idiomatic choice. However, einsum is more flexible. The moment your operation deviates from a simple (...,i,k) @ (...,k,j) pattern—for example, if you need to contract different axes, transpose inputs, or handle higher-order contractions—einsum is the superior, clearer tool. Use matmul for its specific cases and einsum for everything else.

Common Pitfalls

1. Incorrect or Mismatched Subscripts: The most common error is a typo in the subscript string that leads to unintended dimensions. For example, writing 'ik,jk->ij' for matrix multiplication is wrong because the contracted index k must appear in both input matrices. The correct form is 'ik,kj->ij'. Always double-check that the labels for each input array's dimensions match its actual shape and that contraction logic follows the mathematical definition.

1. Ignoring optimize=True for Complex Contractions: Using the default optimize=False for expressions with three or more operands can result in severely suboptimal performance, as it may create massive intermediate arrays. For any non-trivial contraction path, you should habitually set optimize=True. The small overhead of finding the optimal path is almost always worth it.

1. Overusing Einsum for Simple Operations: While einsum can do almost anything, it can be overkill and less readable for very basic operations. Using A.T for transpose, np.trace(A) for the trace, or np.sum(A * B) for an element-wise product followed by sum is often more direct and may be marginally faster. Reserve einsum for operations where its expressiveness truly simplifies the code.

1. Misunderstanding Broadcasting: einsum follows NumPy's broadcasting rules. If you write 'ij,j->ij', it will broadcast the vector j across the rows of the matrix ij. This is powerful but can be a source of error if unintended. Clearly visualize the shape of your output based on the subscript string: axes that appear in an input and the output are free and will be in the result; axes that are only in the inputs are summed.

Summary

• np.einsum uses Einstein summation notation ('subscripts->output') to declaratively specify array operations by summing over repeated indices.
• It elegantly handles matrix multiplication, trace, transpose, diagonal extraction, and, most powerfully, batch operations and tensor contractions in a single, readable line of code.
• For contracting three or more arrays, always use the optimize=True argument to allow NumPy to find the most computationally efficient order of operations, yielding significant speedups.
• While np.matmul is optimal for standard batch matrix multiplication, einsum is the go-to tool for more general or unusual tensor manipulations, such as those frequently required in deep learning gradient computations.
• Avoid common mistakes by carefully aligning subscript labels with array shapes, using optimize=True for complex expressions, and choosing simpler built-in functions for basic tasks like transposition.