Einstein Summation in Python: A Complete Guide to Tensor Operations with NumPy and TensorFlow (2026)

The Einstein Summation notation, often called Einsum, is a compact way to express tensor operations. It is widely used in mathematics, physics, machine learning, and scientific computing because it simplifies complex calculations into readable expressions.

In 2026, Einsum remains an essential tool for machine learning engineers, data scientists, and researchers working with tensor libraries such as NumPy, TensorFlow, and PyTorch. Understanding it helps you write cleaner and more efficient numerical code.

What Is Einstein Summation?

Einstein Summation notation is based on the idea that repeated indices in an expression imply summation. Instead of writing long mathematical formulas with explicit summation symbols, repeated indices automatically indicate which dimensions are summed.

This notation is especially useful for tensor operations, matrix multiplication, vector algebra, and batch computations.

Core Rules of Einstein Summation

1. Repeated indices are summed

If an index appears twice in a term, it is summed over automatically:

$C_{ik} = \sum_{j} A_{ij} B_{jk}$

Here, the index $j$ appears twice, so the operation sums over it.

2. Free indices define the output shape

Indices that appear only once remain in the output tensor. In the equation above, $i$ and $k$ define the shape of the resulting matrix.

Matrix Multiplication With NumPy Einsum

NumPy provides the function np.einsum() for Einstein Summation operations.

import numpy as np

A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])

C = np.einsum('ij,jk->ik', A, B)

print(C)
# [[19 22]
#  [43 50]]

The string ij, jk->ik defines the operation:

• ij corresponds to matrix $A$
• jk corresponds to matrix $B$
• ->ik defines the output shape
• The repeated index $j$ is summed over

Matrix Multiplication With TensorFlow

TensorFlow also supports Einstein Summation through tf.einsum().

import tensorflow as tf

A = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
B = tf.constant([[5, 6], [7, 8]], dtype=tf.float32)

C = tf.einsum('ij,jk->ik', A, B)

print(C)
# tf.Tensor(
# [[19. 22.]
#  [43. 50.]], shape=(2, 2), dtype=float32)

The syntax is identical to NumPy's, making it easy to switch between scientific computing and machine learning workflows.

Common Einsum Examples

Inner product of vectors

The dot product between two vectors is defined as:

$c = \sum_{i} a_i b_i$

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

c = np.einsum('i,i->', a, b)

print(c)
# 32

The repeated index $i$ means multiply element-wise and sum.

Outer product of vectors

The outer product creates a matrix from two vectors:

$C_{ij} = a_i b_j$

C = np.einsum('i,j->ij', a, b)

print(C)
# [[ 4  5  6]
#  [ 8 10 12]
#  [12 15 18]]

No indices are repeated, so no summation occurs.

Transpose of a matrix

You can transpose a matrix by swapping indices:

$A^T_{ji} = A_{ij}$

A_transpose = np.einsum('ij->ji', A)

print(A_transpose)
# [[1 3]
#  [2 4]]

Trace of a matrix

The trace is the sum of diagonal elements:

$\text{Tr}(A) = \sum_{i} A_{ii}$

trace = np.einsum('ii->', A)

print(trace)
# 5

Batch matrix multiplication

Einsum is powerful for batch operations:

$C_{bik} = \sum_{j} A_{bij} B_{bjk}$

A = np.random.rand(3, 2, 2)
B = np.random.rand(3, 2, 2)

C = np.einsum('bij,bjk->bik', A, B)

print(C)

Here, $b$ is the batch dimension, and each matrix pair is multiplied independently.

Advantages of Einsum

Concise notation

Einsum expresses complex tensor operations in a compact form without loops or reshaping.

Flexible operations

It supports matrix multiplication, dot products, transposes, and batch operations in a single interface.

Efficient computation

NumPy and TensorFlow optimize many Einsum operations internally, which can improve performance and memory usage.

When to Use Einsum

Use Einsum when working with:

• Tensor-heavy machine learning models
• Neural network computations
• Batch matrix operations
• Complex linear algebra transformations
• Mathematical formulas that map directly to tensors

For simple operations, standard matrix functions may still be easier to read.

Final Thoughts

Einstein Summation notation is a powerful way to simplify tensor operations. It may look complex at first, but it becomes intuitive with practice.

For developers in 2026 working in machine learning, scientific computing, or AI engineering, Einsum is a valuable tool for writing clean and efficient code.