NumPy - np.einsum() - Einstein Summation

Key Insights

np.einsum() provides a compact notation for expressing complex array operations including matrix multiplication, transpose, trace, and tensor contractions using Einstein summation convention
Understanding the subscript notation is critical: repeated indices indicate summation, while non-repeated indices appear in the output array
For common operations, specialized NumPy functions are often faster, but einsum excels at complex tensor operations and offers superior readability for mathematical expressions

Understanding Einstein Summation Convention

Einstein summation convention eliminates explicit summation symbols by implying summation over repeated indices. In NumPy, np.einsum() implements this convention through a string-based subscript notation that specifies how array dimensions should be combined.

The basic syntax is np.einsum(subscripts, *operands) where subscripts describe the operation using lowercase letters to represent array dimensions. Each letter corresponds to an axis, and the arrow (->) separates input from output specifications.

import numpy as np

# Basic vector dot product: sum over i
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

result = np.einsum('i,i->', a, b)
print(result)  # 32 (1*4 + 2*5 + 3*6)

# Equivalent to
print(np.dot(a, b))  # 32

Matrix Operations

Matrix multiplication demonstrates einsum’s power for linear algebra operations. The subscript notation makes the contraction explicit.

# Matrix multiplication: C[i,j] = sum_k A[i,k] * B[k,j]
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])

C = np.einsum('ik,kj->ij', A, B)
print(C)
# [[19 22]
#  [43 50]]

# Verify with standard matmul
print(np.matmul(A, B))

# Matrix-vector multiplication
v = np.array([1, 2])
result = np.einsum('ij,j->i', A, v)
print(result)  # [5 11]

# Batch matrix multiplication
batch_A = np.random.rand(10, 3, 4)
batch_B = np.random.rand(10, 4, 5)
batch_C = np.einsum('bij,bjk->bik', batch_A, batch_B)
print(batch_C.shape)  # (10, 3, 5)

Transpose and Trace Operations

Einsum handles array rearrangement elegantly without requiring separate function calls.

# Matrix transpose
A = np.array([[1, 2, 3], [4, 5, 6]])
A_T = np.einsum('ij->ji', A)
print(A_T)
# [[1 4]
#  [2 5]
#  [3 6]]

# Trace (sum of diagonal elements)
M = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
trace = np.einsum('ii->', M)
print(trace)  # 15 (1 + 5 + 9)

# Diagonal extraction
diag = np.einsum('ii->i', M)
print(diag)  # [1 5 9]

# Multi-dimensional transpose
tensor = np.random.rand(2, 3, 4, 5)
transposed = np.einsum('ijkl->kjil', tensor)
print(transposed.shape)  # (4, 3, 2, 5)

Element-wise and Broadcasting Operations

Einsum supports element-wise operations and broadcasting through its subscript notation.

# Element-wise multiplication
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
result = np.einsum('ij,ij->ij', A, B)
print(result)
# [[ 5 12]
#  [21 32]]

# Outer product
a = np.array([1, 2, 3])
b = np.array([4, 5])
outer = np.einsum('i,j->ij', a, b)
print(outer)
# [[ 4  5]
#  [ 8 10]
#  [12 15]]

# Column-wise sum
A = np.array([[1, 2, 3], [4, 5, 6]])
col_sum = np.einsum('ij->j', A)
print(col_sum)  # [5 7 9]

# Row-wise sum
row_sum = np.einsum('ij->i', A)
print(row_sum)  # [ 6 15]

Advanced Tensor Contractions

Einsum truly shines when working with higher-dimensional tensors where standard operations become cumbersome.

# Bilinear form: x^T A y
A = np.array([[1, 2], [3, 4]])
x = np.array([1, 2])
y = np.array([3, 4])

result = np.einsum('i,ij,j->', x, A, y)
print(result)  # 50

# Tensor contraction over multiple indices
T1 = np.random.rand(3, 4, 5)
T2 = np.random.rand(4, 5, 6)
result = np.einsum('ijk,jkl->il', T1, T2)
print(result.shape)  # (3, 6)

# Batched dot products
vectors_a = np.random.rand(100, 3)
vectors_b = np.random.rand(100, 3)
dot_products = np.einsum('bi,bi->b', vectors_a, vectors_b)
print(dot_products.shape)  # (100,)

Practical Applications

Here are real-world scenarios where einsum provides clean, efficient solutions.

# Attention mechanism (simplified)
queries = np.random.rand(32, 10, 64)   # (batch, seq_len, dim)
keys = np.random.rand(32, 10, 64)
values = np.random.rand(32, 10, 64)

# Compute attention scores: Q @ K^T
scores = np.einsum('bqd,bkd->bqk', queries, keys)
print(scores.shape)  # (32, 10, 10)

# Apply attention to values
attention_weights = np.random.rand(32, 10, 10)  # After softmax
output = np.einsum('bqk,bkd->bqd', attention_weights, values)
print(output.shape)  # (32, 10, 64)

# Covariance matrix from centered data
X = np.random.rand(1000, 50)  # 1000 samples, 50 features
X_centered = X - X.mean(axis=0)
cov = np.einsum('ij,ik->jk', X_centered, X_centered) / (X.shape[0] - 1)
print(cov.shape)  # (50, 50)

# Weighted sum along specific axes
data = np.random.rand(10, 20, 30)
weights = np.random.rand(20)
weighted = np.einsum('ijk,j->ik', data, weights)
print(weighted.shape)  # (10, 30)

Performance Considerations

While einsum is elegant, performance varies based on operation complexity and array sizes.

import time

# Compare performance
A = np.random.rand(500, 500)
B = np.random.rand(500, 500)

# Using einsum
start = time.time()
for _ in range(100):
    C1 = np.einsum('ij,jk->ik', A, B)
einsum_time = time.time() - start

# Using matmul
start = time.time()
for _ in range(100):
    C2 = np.matmul(A, B)
matmul_time = time.time() - start

print(f"einsum: {einsum_time:.4f}s")
print(f"matmul: {matmul_time:.4f}s")

# Enable optimization with optimize parameter
start = time.time()
for _ in range(100):
    C3 = np.einsum('ij,jk->ik', A, B, optimize=True)
optimized_time = time.time() - start

print(f"einsum optimized: {optimized_time:.4f}s")

The optimize parameter enables path optimization for complex contractions with multiple operands, potentially providing significant speedups.

# Complex multi-tensor contraction
T1 = np.random.rand(10, 20, 30)
T2 = np.random.rand(30, 40, 50)
T3 = np.random.rand(50, 60)

# Without optimization
result1 = np.einsum('ijk,klm,mn->ijn', T1, T2, T3, optimize=False)

# With optimization (finds efficient contraction order)
result2 = np.einsum('ijk,klm,mn->ijn', T1, T2, T3, optimize=True)

# Results are identical
print(np.allclose(result1, result2))  # True

Common Patterns and Idioms

Recognize these patterns to quickly translate mathematical operations into einsum notation.

# Sum all elements
A = np.random.rand(3, 4, 5)
total = np.einsum('ijk->', A)

# Sum over specific axes (axis=1)
sum_axis1 = np.einsum('ijk->ik', A)

# Kronecker product
a = np.array([1, 2])
b = np.array([3, 4, 5])
kron = np.einsum('i,j->ij', a, b).ravel()

# Hadamard product (element-wise) followed by sum
A = np.random.rand(100, 100)
B = np.random.rand(100, 100)
hadamard_sum = np.einsum('ij,ij->', A, B)

# Frobenius inner product
frob = np.einsum('ij,ij->', A, B)
print(f"Frobenius: {frob}")

Einstein summation through np.einsum() transforms complex tensor operations into readable, maintainable code. Master the subscript notation and you’ll handle multidimensional array operations with confidence and clarity.