PyPIAI / MLupdated 08 Aug 2026

opt-einsum

Item: opt-einsum
Rating: 4.8
Author: Keyoor

opt-einsum finds an efficient pairwise execution order for Einstein-summation tensor expressions and then runs the contraction on NumPy, PyTorch, JAX, TensorFlow, Dask, CuPy, or another supported array backend. The same mathematical expression can differ by orders of magnitude in work and intermediate memory depending on which tensors are combined first. This package optimizes that path; it does not invent the equation, provide tensors, or guarantee the globally best path cheaply.

opt-einsum · v3.4.0 · 5.9M/wk · ★ 989 · pushed 19 Mar 2026 · dgasmith/opt_einsum · docs · reviewed by Keyoor

Verdict

opt-einsum earns its place when a real multi-tensor contraction is expensive enough to inspect, cache, and benchmark. Skip it for ordinary matrix math or when a whole-graph compiler should own execution planning.

API stability5/5contract, contract_path, contract_expression, backend dispatch, optimizer names, memory_limit, and shared_intermediates form a mature compact API. Version 3.4.0 retains NumPy-style subscript semantics and accepts explicit path lists, so optimized plans can be stored independently of a particular search strategy. Most recent development is additive backend and typing work rather than conceptual churn.

Docs5/5The documentation explains path introduction, optimal and heuristic strategies, reusable paths, contract expressions, constants, shared intermediates, backends, large expressions, custom optimizers, and API signatures with cost tables. It states the search-complexity and memory-limit tradeoffs directly, which is essential because a faster arithmetic count does not automatically mean a faster workload.

Maintenance4/5The repository had 989 stars, 37 open issues and PRs, and a push on 2026-03-19. PyPI 3.4.0 was released on 2024-09-26 and supports Python through 3.13 in its classifiers. The release cadence is measured rather than frequent, but the narrow algorithmic core is mature and the repository continues receiving maintenance.

Ecosystem5/5The README documents NumPy, Dask, PyTorch, TensorFlow, JAX, CuPy and other backends, and part of the project's optimization work has been incorporated upstream into NumPy. The backend-neutral interface and explicit path format fit scientific Python and machine-learning code well, while external optimizers such as cotengra can address larger tensor-network searches.

Use it if

Your einsum expression has three or more operands and contraction order materially affects runtime or memory
You repeat a fixed tensor equation and want to precompute and reuse its path or constant subexpressions
You need one contraction interface across NumPy, PyTorch, JAX, TensorFlow, Dask, or GPU array libraries
You want path diagnostics that report scaling, estimated operations, and largest intermediate before running expensive work

Skip it if

Your expressions are simple matrix products or two-operand einsums: NumPy and tensor frameworks already route common cases to tuned kernels, so path optimization adds little
You expect optimal search to be cheap for many operands: the documentation states the optimal algorithm scales factorially, while branch and dynamic-programming searches can also become expensive
Peak memory must obey a byte-level device budget: memory_limit is expressed in tensor elements and can force much slower contraction paths; dtype and backend allocations still need separate accounting
You need readable shape manipulation rather than contraction speed: einops names and rearranges axes more clearly, while opt-einsum keeps single-character einsum notation
Your framework compiler already optimizes and fuses the complete model graph: a locally chosen pairwise path can conflict with broader compilation, sharding, or memory-planning decisions

Setup reality

pip install opt-einsum installs a pure-Python package with no declared runtime dependencies; 3.4.0 requires Python 3.8 or newer. It does not install NumPy, PyTorch, JAX, TensorFlow, CuPy, Dask, or a GPU runtime. Backend selection normally follows the first operand type, so the compatible array framework must already be present and correctly configured. The first surprise is that optimization itself costs time. auto aims to keep path search under about one millisecond, auto-hq spends up to roughly a second for a better path, greedy is cheap, and optimal is factorial in operand count. Benchmark the search and execution together for one-off contractions, then use contract_expression or a saved path for repeated shapes. Paths depend on dimension sizes, not just the subscript string; reusing one for very different shapes can be valid but poor. The optimizer estimates operations and intermediate element counts, but real performance also depends on BLAS, device transfers, dtype, tensor layout, framework dispatch, compilation, and available memory. memory_limit can cap the largest intermediate by element count or max_input, but the docs warn that a limit may make execution exponentially slower. Backend conversion through a ContractExpression can move NumPy arrays to another backend and back, which is convenient but can erase any compute win through transfer overhead. Shared intermediates retain arrays until the context and cache are released, so it trades memory for repeated work. Finally, recent NumPy already incorporates much of opt-einsum's path work; confirm this separate dependency adds features or backend behavior you actually use.

Patterns

Run an optimized Einstein contractioncontract-tensors

import numpy as np
import opt_einsum as oe

a = np.random.rand(20, 30)
b = np.random.rand(30, 40)
c = oe.contract('ij,jk->ik', a, b)

Two-operand matrix multiplication is already a strong backend case; the larger gains appear when three or more tensors permit different paths.

Use cheap greedy path searchchoose-greedy-path

result = oe.contract(
    'ab,bc,cd,de->ae', a, b, c, d,
    optimize='greedy',
)

Greedy search scales well for many operands but is heuristic; compare it with auto-hq when the contraction is repeated enough to justify more search time.

Inspect work and peak-intermediate estimatesinspect-contraction-path

path, info = oe.contract_path(
    'ab,bc,cd->ad', a, b, c,
    optimize='auto-hq',
)
print(path)
print(info)

PathInfo reports estimated arithmetic and element counts, not measured device time or byte-accurate allocator usage.

Find a path without allocating tensorsplan-from-shapes

path, info = oe.contract_path(
    'ab,bc,cd->ad',
    (1000, 64), (64, 64), (64, 1000),
    shapes=True, optimize='greedy',
)

With shapes=True every operand must be a shape tuple; mixing arrays and shapes produces misleading parsing errors.

Reuse a path across callsreuse-explicit-path

path, _ = oe.contract_path('ab,bc,cd->ad', a, b, c, optimize='auto-hq')

for a_batch in batches:
    result = oe.contract('ab,bc,cd->ad', a_batch, b, c, optimize=path)

A path remains valid for compatible ranks and dimensions, but it may cease to be efficient when shape sizes change materially.

Build a callable contraction for fixed shapescompile-reusable-expression

expr = oe.contract_expression(
    'ab,bc,cd->ad',
    (100, 64), (64, 64), (64, 20),
    optimize='auto-hq',
)
result = expr(a, b, c)

The expression is planned for the supplied shapes; it accepts compatible ranks but a changed size profile can make the cached plan suboptimal.

Cache work involving constant operandsprecompute-constant-tensors

expr = oe.contract_expression(
    'ab,bc,cd->ad',
    a, (64, 64), c,
    constants=[0, 2], optimize='auto-hq',
)
result = expr(b)

Pass actual arrays at constant positions and shapes elsewhere; constants are retained and evaluated per backend, increasing lifetime and memory use.

Dispatch directly on PyTorch tensorsuse-torch-backend

import torch
import opt_einsum as oe

a = torch.randn(32, 64, device='cuda')
b = torch.randn(64, 16, device='cuda')
result = oe.contract('ab,bc->ac', a, b, backend='torch')

opt-einsum does not install PyTorch or configure CUDA; operands and the backend runtime must already be on the intended device.

Reuse a preallocated output arraywrite-into-output

out = np.empty((a.shape[0], c.shape[1]))
oe.contract('ab,bc->ac', a, c, out=out)
assert out.shape == (a.shape[0], c.shape[1])

The output shape and dtype must be compatible; intermediate allocations can still occur even when the final array is supplied.

Cap the largest planned intermediatelimit-intermediate-size

result = oe.contract(
    equation, *operands,
    optimize='auto-hq',
    memory_limit='max_input',
)

The limit counts elements rather than bytes and can force dramatically more arithmetic, so measure both memory and runtime.

Reuse common work across contractionsshare-intermediate-results

with oe.shared_intermediates() as cache:
    left = oe.contract('ab,bc,cd->ad', a, b, c)
    right = oe.contract('ab,bc,ce->ae', a, b, e)
print(f'cached operations: {len(cache)}')

The cache holds references to tensors and intermediates until it is released; sharing saves compute by deliberately consuming memory.

Avoid the single-character subscript limituse-integer-index-labels

result = oe.contract(
    a, [0, 1],
    b, [1, 2],
    [0, 2],
)

Interleaved notation alternates each operand with its index list and ends with output indices; it is useful for generated high-rank equations.

Alternatives

Package	Registry	Pick it when
numpy	PyPI	Your contractions stay in NumPy and its built-in einsum optimize modes are sufficient
cotengra	PyPI	You need more advanced hypergraph contraction optimization for large tensor networks
einops	PyPI	Your main problem is readable reshaping, axis rearrangement, repetition, and reductions
tensornetwork	PyPI	You want an explicit tensor-network graph abstraction in addition to contraction