opt-einsum review
opt-einsum 3.4.0 chooses an execution order for Einstein summations, then dispatches the pairwise contractions to the array backend already holding your tensors. Changing that order can reduce arithmetic and the largest intermediate without changing the equation. Our import took 0.12 seconds, and the pure-Python package installed with no dependencies. Version 3.4.0 removed NumPy as a required dependency, added type annotations, moved the docs to MkDocs, accepted `backend=None`, added a `jaxlib` alias, and fixed dynamic-programming path failures involving very low memory limits and scalar-only contractions.
opt-einsum 3.4.0 installed in 0.3 seconds as one 1 MB package and imported in 0.12 seconds in our sandbox, so path planning adds almost no installation baggage to an existing tensor stack. Use it for repeated 3-plus-operand contractions whose path and intermediate size you can measure; leave 2-operand math and whole-graph compiler scheduling to the backend.
We installed it
| Install | ✓ · 0.3s | 1 package on disk · 1 MB |
| Import | ✓ | import opt_einsum in 0.12s · pure Python · requires Python >=3.8 |
| Known vulns | 0 | (pip-audit) |
Answers from our run
Does opt-einsum install cleanly?
Yes. In a fresh container with an empty cache, pip install opt-einsum finished in 0.3s, leaving 1 package and 1 MB on disk. pip-audit reported no known vulnerabilities.
What does opt-einsum need to run?
Python >=3.8, and nothing compiled: it is pure Python. In our run import opt_einsum succeeded in 0.12s.
opt-einsum or numpy: which should you use?
numpy: Use its built-in optimized einsum when every operand is already a NumPy array. opt-einsum 3.4.0 installed in 0.3 seconds as one 1 MB package and imported in 0.12 seconds in our sandbox, so path planning adds almost no installation baggage to an existing tensor stack.
When should you not use opt-einsum?
The work is ordinary matrix multiplication or a 2-operand einsum. There is no meaningful multi-step path to search.
Use it if
- An einsum has at least 3 operands and several legal pairwise orders with different intermediate sizes.
- The same equation and shapes repeat often enough to cache a path or a `ContractExpression`.
- One contraction API must dispatch across NumPy, PyTorch, JAX, TensorFlow, Dask, or CuPy arrays.
- You need estimated operation counts and largest-intermediate figures before committing to an expensive run.
- The work is ordinary matrix multiplication or a 2-operand einsum. There is no meaningful multi-step path to search.
- You plan to run `optimal` search across many tensors. The documentation says that strategy scales factorially with operand count.
- A device must stay below an exact byte limit. `memory_limit` counts tensor elements, and backend allocations still depend on dtype and runtime behavior.
- The task is axis rearrangement or readable reshaping. `einops` describes those transformations directly instead of using contraction notation.
- A framework compiler owns whole-model fusion, sharding, and memory scheduling. A locally fixed pairwise order may work against that wider plan.
Setup reality
We installed opt-einsum 3.4.0 in 0.3 seconds with Python 3.12 in a fresh Bookworm sandbox. It left 1 package and 1 MB on disk. The wheel is pure Python, declares 0 direct dependencies, requires Python 3.8 or newer, and carries the MIT License. import opt_einsum worked in 0.12 seconds, while pip-audit found 0 known vulnerabilities. The distribution does not include a py.typed marker.
Version 3.4.0 no longer installs NumPy. You must supply NumPy, PyTorch, JAX, TensorFlow, CuPy, Dask, or another supported array system, plus any CUDA runtime and device configuration it needs. Backend inference normally follows operand types. No credentials or config file are required by opt-einsum itself.
Path search has its own cost. greedy is the cheap heuristic; optimal grows factorially as operands are added. Save the returned path or create a ContractExpression when shapes repeat. A path remains executable for compatible inputs, but changed dimension sizes can make its earlier ordering a poor choice.
memory_limit constrains intermediate element counts, not bytes, and the documentation warns that tight limits can greatly increase arithmetic. shared_intermediates() keeps tensor references alive until its cache is released. Backend conversion can also move arrays between host and accelerator, so operation estimates alone do not include transfer time, allocator behavior, dtype width, or compiler fusion.
Patterns
Contract two NumPy arrays contract-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)With 2 operands there is no multi-step order to optimize; use this form mainly for API consistency across larger equations.
Pick the greedy optimizer choose-greedy-path
result = oe.contract(
'ab,bc,cd,de->ae', a, b, c, d,
optimize='greedy',
)`greedy` searches quickly and may miss a cheaper path. Compare `auto-hq` when the same 4-operand equation runs many times.
Print the planned contraction steps inspect-contraction-path
path, info = oe.contract_path(
'ab,bc,cd->ad', a, b, c,
optimize='auto-hq',
)
print(path)
print(info)`PathInfo` contains estimated operations and element counts. It does not measure accelerator time or allocator bytes.
Plan from shape tuples plan-from-shapes
path, info = oe.contract_path(
'ab,bc,cd->ad',
(1000, 64), (64, 64), (64, 1000),
shapes=True, optimize='greedy',
)When `shapes=True`, pass a tuple for every operand. Do not mix live arrays into the same planning call.
Apply one saved path repeatedly reuse-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)Reusing the list avoids another search. Recompute it when dimension sizes change enough to alter intermediate costs.
Compile a shape-specific expression compile-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 callable is planned against these 3 shape tuples; a different size profile can keep working while running inefficiently.
Precompute constant sides of an expression precompute-constant-tensors
expr = oe.contract_expression(
'ab,bc,cd->ad',
a, (64, 64), c,
constants=[0, 2], optimize='auto-hq',
)
result = expr(b)Positions 0 and 2 receive arrays because they are constants. The expression retains them, which extends their memory lifetime.
Run on CUDA tensors through PyTorch use-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')The opt-einsum wheel has 0 dependencies, so PyTorch and a working CUDA runtime must already be installed.
Write the final result into an array write-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])`out` avoids allocating the final result only. Pairwise steps can still allocate intermediates.
Limit intermediate element counts limit-intermediate-size
result = oe.contract(
equation, *operands,
optimize='auto-hq',
memory_limit='max_input',
)`max_input` is measured in elements rather than bytes, and the selected path may trade much more arithmetic for that cap.
Share repeated pairwise results share-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 context keeps intermediate tensors referenced until exit. Inspect the cache length and peak memory before using this in a long batch.
Use integer labels for generated equations use-integer-index-labels
result = oe.contract(
a, [0, 1],
b, [1, 2],
[0, 2],
)Interleaved syntax pairs each operand with an index list and puts output indexes last, avoiding single-letter label limits.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| numpy | PyPI | Use its built-in optimized einsum when every operand is already a NumPy array. |
| cotengra | PyPI | Use it for hypergraph path search on large tensor networks. |
| einops | PyPI | Use it when named reshaping and axis rearrangement are the actual problem. |
| tensornetwork | PyPI | Use it when the code should model a tensor network explicitly as well as contract it. |
More ai / ml guides
openai · mcp · huggingface-hub · scikit-learn · tiktoken · @modelcontextprotocol/sdk · the whole shelf →
How this guide is made: grounded in the library's documentation, release notes, changelog, and issue history, on a fixed rubric — not a hands-on install of every release. The 50 most-downloaded entries are additionally install-verified in clean containers. Corrections: contact the desk.