Einsum Patterns

Use this page to understand fnp.einsum -- the core computation primitive in flopscope.

You will learn:

How to write common einsum patterns and understand their FLOP costs
How to use symmetric tensors with einsum for cost savings
How to inspect and customize contraction paths
How to leverage path caching for repeated operations

Common patterns

import flopscope as flops
import flopscope.numpy as fnp

with flops.BudgetContext(flop_budget=10**8) as budget:
    A = fnp.ones((256, 256))
    B = fnp.ones((256, 256))
    x = fnp.ones((256,))

    # Matrix-vector multiply
    y = fnp.einsum('ij,j->i', A, x)           # 131,072 FLOPs

    # Matrix multiply
    C = fnp.einsum('ij,jk->ik', A, B)         # 33,554,432 FLOPs

    # Outer product
    outer = fnp.einsum('i,j->ij', x, x)       # 65,792 FLOPs

    # Trace
    tr = fnp.einsum('ii->', A)                 # 256 FLOPs

    # Batched matmul
    batch = fnp.ones((4, 256, 256))
    out = fnp.einsum('bij,bjk->bik', batch, batch)  # 134,217,728 FLOPs

    print(budget.summary())

The cost of an einsum is the sum of per-step costs along the optimal contraction path. Every einsum — even a simple two-operand one — goes through the opt_einsum path optimizer (a symmetry-aware fork of opt_einsum).

For each pairwise contraction step the cost decomposes into a multiplicative term μ and an accumulation term α:

cost = μ + α

where μ is the product of all index dimensions (the FMAs you'd issue) and α charges the accumulation across the contracted axes. The α/M model counts multiplies and adds separately (FMA=2 textbook); the factor-of-2 you see for dense matmul comes from μ counting multiplies and α counting accumulation adds.

For 'ij,jk->ik' with shapes (256, 256) and (256, 256):

Indices: i=256, j=256, k=256
μ = 256 × 256 × 256 = 16,777,216
α = 16,777,216
Cost: 33,554,432

For multi-operand einsums (3+ tensors), Flopscope automatically decomposes the contraction into optimal pairwise steps. The total cost is the sum of per-step costs.

When symmetric tensors are involved, each step's cost is further reduced by the ratio of unique output elements to total output elements. See Symmetry Savings for the full practical guide.

fnp.dot and fnp.matmul

fnp.dot(A, B) and fnp.matmul(A, B) are equivalent to the corresponding einsum and have the same FLOP cost.

Symmetric tensors

There are two separate symmetry declarations — one for inputs, one for outputs:

Input symmetry — wrap with flops.as_symmetric() before passing to einsum. The optimizer automatically uses symmetry to choose the best contraction order and charges reduced costs:

with flops.BudgetContext(flop_budget=10**8) as budget:
    S = flops.as_symmetric(fnp.eye(10), symmetry=flops.SymmetryGroup.symmetric(axes=(0, 1)))  # 55 unique elements
    v = fnp.ones((10,))

    result = fnp.einsum('ij,j->i', S, v)  # cost reduced by input symmetry

Output symmetry — pass symmetry= to einsum() to declare that the result is symmetric. This wraps the output as a SymmetricTensor so downstream operations benefit from reduced costs. It does NOT affect the cost of this einsum — it's a declaration about the result's structure:

with flops.BudgetContext(flop_budget=10**8) as budget:
    X = fnp.random.randn(100, 10)

    # X^T X is always symmetric — declare the exact output group
    C = fnp.einsum('ki,kj->ij', X, X, symmetry=flops.SymmetryGroup.symmetric(axes=(0, 1)))

    print(type(C))  # <class 'SymmetricTensor'>
    # C can now be passed to other operations with automatic cost savings

For the full symmetry guide, see Symmetry Savings.

Inspecting costs

fnp.einsum_path() previews the contraction plan without executing the contraction itself. Planning is cheap: it records a nominal 1-FLOP einsum_path event, but none of the contraction FLOPs are spent.

import flopscope as flops
import flopscope.numpy as fnp

n = 10
T = flops.as_symmetric(fnp.ones((n, n, n)), symmetry=flops.SymmetryGroup.symmetric(axes=(0, 1, 2)))
A = fnp.random.randn(n, n)
B = fnp.random.randn(n, n)
C = fnp.random.randn(n, n)

path, info = fnp.einsum_path('ijk,ai,bj,ck->abc', T, A, B, C)

print(f"Path: {path}")
print(info)
print(f"Optimized cost: {info.optimized_cost:,}")
print(f"Optimizer used: {info.optimizer_used}")

Path: [(0, 1), (0, 2), (0, 1)]
  Complete contraction:  ijk,ai,bj,ck->abc
      Naive cost (flopscope):  4,000,000
  Optimized cost (flopscope):  4,000,000
                     Speedup:  1.000x
                     Savings:  0.0%
       Largest intermediate:  1,000 elements
                Index sizes:  a=b=c=i=j=k=10
                  Optimizer:  optimal
------------------------------------------------------------------------------------
step  contract  subscript                                flops  blas
------------------------------------------------------------------------------------
   0  (0, 1)    ai,ijk->ajk                             10,000  GEMM
   1  (0, 2)    ajk,bj->akb                             10,000  TDOT
   2  (0, 1)    akb,ck->abc                             10,000  TDOT
Optimized cost: 4,000,000
Optimizer used: optimal

The printed table shows the contraction order, the α/M cost (path-independent — naive and optimized are equal by construction), largest intermediate, grouped index sizes, and one row per pairwise contraction step. Each step shows the chosen contract tuple, subscript, per-step FLOPs, and BLAS tag.

For per-step debugging, call print(info.format_table(verbose=True)). The verbose view adds indented rows with the merged operand subset, the intermediate output shape, and the running cumulative cost.

For the symmetric-vs-dense FLOP breakdown that the path-independent total doesn't expose, use flops.einsum_accumulation_cost(...).describe() — see the inspection section below.

Inspecting what gets charged

If you want to see the exact FLOP cost flopscope will charge for an einsum (without actually running it), call flops.einsum_accumulation_cost:

import flopscope as flops
import flopscope.numpy as fnp

A = fnp.zeros((4, 4))
B = fnp.zeros((4, 4))

cost = flops.einsum_accumulation_cost('ij,jk->ik', A, B)
print(cost.total)  # 128

The cost is path-independent: flops.einsum_accumulation_cost(...).total matches info.optimized_cost from fnp.einsum_path(...), and both equal the FLOPs that BudgetContext will charge when you actually call fnp.einsum(...). Repeated calls hit an LRU cache and return in microseconds — see flops.einsum_cache_info() and flops.einsum_clear_caches() if you need to reset between cold-call benchmarks.

For inputs with declared symmetry, the returned total reflects the orbit-aware savings — see Symmetry-aware FLOP counting for the mental model.

Custom contraction paths

By default Flopscope finds the optimal contraction order automatically. You can override this by passing an explicit path — a list of int-tuples specifying which operand positions to contract at each step:

import flopscope as flops
import flopscope.numpy as fnp

A = fnp.ones((3, 4))
B = fnp.ones((4, 5))
C = fnp.ones((5, 6))

# Plan first, execute later
path, info = fnp.einsum_path('ij,jk,kl->il', A, B, C)
print(f"Optimal path: {path}")  # e.g. [(0, 1), (0, 1)]

# Execute with the planned path
with flops.BudgetContext(flop_budget=10**8) as budget:
    result = fnp.einsum('ij,jk,kl->il', A, B, C, optimize=path)

You can also specify a completely custom path. Each tuple names the positions (in the current operand list) to contract; the result is appended to the end:

# Force B×C first (positions 1,2), then A×result (positions 0,1)
result = fnp.einsum('ij,jk,kl->il', A, B, C, optimize=[(1, 2), (0, 1)])

# Force A×B first (positions 0,1), then result×C (positions 0,1)
result = fnp.einsum('ij,jk,kl->il', A, B, C, optimize=[(0, 1), (0, 1)])

Different paths may have different FLOP costs. Use fnp.einsum_path() to compare — it returns the plan without executing the contraction.

Path caching

Contraction paths are cached automatically in a module-level LRU cache. When you call fnp.einsum() with the same subscripts, shapes, optimizer, and symmetry structure, the path is reused from cache instead of being recomputed. This makes repeated einsums in loops essentially free in path-finding overhead:

with flops.BudgetContext(flop_budget=10**9) as budget:
    for i in range(1000):
        y = fnp.einsum('ij,j->i', A, x)  # path computed once, reused 999 times

fnp.einsum_path() shares the same cache, so planning a path warms the cache for subsequent fnp.einsum() calls and vice versa.

Cache management

# Inspect cache statistics
info = fnp.einsum_cache_info()
print(f"Hits: {info.hits}, Misses: {info.misses}, Size: {info.currsize}/{info.maxsize}")

# Clear the cache (e.g., to free memory or force recomputation)
fnp.clear_einsum_cache()

# Change the cache size (default 4096 entries, rebuilds the cache)
flops.configure(einsum_path_cache_size=8192)

Common pitfalls

Symptom: Unexpectedly high FLOP cost

Fix: Check all index dimensions. A subscript like 'ijkl,jklm->im' multiplies all five dimension sizes together. Use flops.accounting.einsum_cost() or fnp.einsum_path() to preview costs before executing.

Symmetry Savings — full guide to symmetry mechanisms
API Reference — algorithms, symmetry support, and operation details
Plan Your Budget — query costs before executing
FLOP Counting Model — how costs are computed