Registry / ai-ml / opt-einsum

opt-einsum

JSON →
library3.4.0pypypi✓ verified 52d ago

opt-einsum is a Python library that optimizes the contraction order of Einstein summation expressions, significantly reducing the execution time of einsum-like operations in various backends such as NumPy, Dask, PyTorch, TensorFlow, and JAX. It achieves this by finding efficient contraction paths, often dispatching operations to highly optimized routines like BLAS or cuBLAS. The library is currently at version 3.4.0 and is actively maintained, serving as the underlying optimization engine for `numpy.einsum(..., optimize=True)` and `torch.einsum` when installed.

ai-mldatahttp-networking
pip install opt-einsum
Install & Compatibility
Where this runs
tested against v3.4.0 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
musl
py 3.103.925 runs
installs and imports cleanly · install 0.0s · import 0.038s · 18.3MB
glibc
py 3.103.925 runs
installs and imports cleanly · install 1.6s · import 0.031s · 19MB
16MB installed
● package 16MB
Code
Verified usage

Verified import paths — ran on the pinned version, not inferred.

contract
from opt_einsum import contract
import opt_einsum as oe; oe.einsum(...)
The primary function `contract` is a direct drop-in replacement for `np.einsum` with optimization.
contract_path
from opt_einsum import contract_path
Used to inspect and pre-calculate the optimal contraction path for reuse.

This quickstart demonstrates how to use `opt_einsum.contract` as a drop-in replacement for `numpy.einsum` to automatically optimize the tensor contraction order and achieve significant performance improvements.

import numpy as np from opt_einsum import contract N = 10 C = np.random.rand(N, N) I = np.random.rand(N, N, N, N) # Using unoptimized numpy.einsum (for comparison) # result_np = np.einsum('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C) # Using opt_einsum.contract for optimized performance result_opt = contract('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C) print(f"Optimized result shape: {result_opt.shape}")
Debug
Known issues
breakingThe `path` keyword argument in `opt_einsum.contract` has been changed to `optimize` to align more closely with NumPy's API. The `path` keyword will be deprecated in future versions.
fix
Replace `path=...` with `optimize=...` in calls to `opt_einsum.contract`.
affects: >=3.4.0
gotchaWhen using `numpy.einsum`, explicitly setting `optimize=True` (or a specific path strategy) is crucial for performance. Without `optimize`, `numpy.einsum` defaults to a left-to-right contraction order, which can be highly inefficient for complex expressions. `opt-einsum.contract` applies optimization by default.
fix
Always use `np.einsum(..., optimize=True)` or `opt_einsum.contract(...)` for performance-critical einsum operations.
affects: All versions of NumPy and opt-einsum
gotchaFinding the truly optimal contraction path for an einsum expression is an NP-hard problem. While `opt-einsum` offers an 'optimal' strategy, it can scale factorially with the number of terms and quickly become intractable for many tensors. For larger expressions, heuristic algorithms like 'greedy' or 'random-greedy' are used.
fix
For complex or many-tensor contractions, prefer `optimize='auto'` (the default) or explicitly choose a heuristic path like `'greedy'` or `'random-greedy-128'` to balance path quality and computation time. Avoid `'optimal'` for expressions with more than a few tensors.
affects: All versions
gotchaThe `memory_limit` parameter in `opt_einsum.contract` can constrain the size of intermediate tensors. While useful for memory management, imposing a limit can make contractions exponentially slower to perform if it restricts the optimizer from finding the most efficient path. The default is `None`, meaning no memory limit.
fix
Carefully consider the trade-off between memory usage and performance when setting `memory_limit`. Only use it if memory constraints are strict, and be aware of potential performance degradation.
affects: All versions
breaking`numpy` is a required dependency for `opt-einsum` and must be installed for `opt-einsum` functionality (and any `numpy.einsum` usage). The `ModuleNotFoundError` indicates `numpy` was not found in the environment.
fix
Ensure `numpy` is installed in the environment (e.g., `pip install numpy`) before using `opt-einsum` or any related functionality.
affects: All versions
breakingThe `opt-einsum` library requires `numpy` as a core dependency. A `ModuleNotFoundError` indicates that `numpy` is not installed in the environment, preventing the library from functioning.
fix
Ensure `numpy` is installed in your Python environment, typically by running `pip install numpy`.
affects: All versions
Errors
Common errors & fixes
ValueError: invalid subscript 'Ų' in einstein sum subscripts string, subscripts must be letters
The optimized contraction path, especially when `memory_limit` is active or with very complex expressions, generates an intermediate `einsum` equation with more unique indices than available single lowercase letters (a-z), which some backends and `opt_einsum` itself enforce.
fix
Simplify the `einsum` expression, break it into smaller parts, or explore different optimization strategies (e.g., `optimize='greedy'`) if the default or `memory_limit` causes this. If using `contract_path`, inspect the `path_info` to understand the intermediate expression.
ImportError: cannot import name 'tensorflow' from 'opt_einsum.backends'
This error occurs when `opt_einsum` attempts to load its TensorFlow backend but cannot find the `tensorflow` module as expected, often due to an incomplete or corrupted `opt_einsum` installation or an environment issue.
fix
Uninstall and reinstall `opt_einsum` (`pip uninstall opt_einsum && pip install opt_einsum`). Ensure TensorFlow is correctly installed and accessible in the Python environment.
AttributeError: module 'torch.backends' has no attribute 'opt_einsum'
In some newer PyTorch versions, `torch.backends.opt_einsum` might not be directly exposed as a public attribute, even though `opt_einsum` can still be used as a backend for `torch.einsum` if installed.
fix
To explicitly interact with the `opt_einsum` backend for PyTorch, use `import torch.backends.opt_einsum` directly. Alternatively, simply ensure `opt-einsum` is installed (`pip install opt-einsum`) for `torch.einsum` to leverage it automatically for optimization.
ValueError: axes don't match array
This error typically occurs when `numpy.einsum` is called with `optimize=True`, which often delegates to `opt_einsum`'s path-finding algorithms. If the original `einsum` expression contains implicit summations or index patterns that the optimized path (which might use `tensordot`) cannot correctly interpret due to mismatching or implicitly summed axes, this error can arise.
np.einsum('bdc,ac->ab', a, b, optimize=True) fails but works with optimize=False
When `numpy.einsum` with `optimize=True` tries to find an optimal contraction path, it might use intermediate operations like `tensordot` that do not inherently support certain implicit summations or index patterns as flexibly as the default `einsum` implementation, leading to axis mismatch errors.
Upgrade
Version history
3.4.0latest on PyPI
Audit
Dependencies
numpyrequiredCommonly used for array operations and a typical backend for einsum expressions.
torchoptionalUsed as a backend for PyTorch tensors, offering optimized einsum for PyTorch.
tensorflowoptionalUsed as a backend for TensorFlow tensors, offering optimized einsum.
daskoptionalUsed as a backend for Dask arrays, enabling optimized einsum on larger-than-memory or distributed arrays.
cupyoptionalUsed as a backend for CuPy arrays, enabling optimized einsum on GPUs.
Agent activity
18 hits · last 30 days
seranking-bot
4
amazonbot
3
node
2
ahrefsbot
2
Meta
1
Resources