opt-einsum

3.4.0 · active · verified Sun Mar 29

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.

Warnings

breaking The `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`.
gotcha When 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.
gotcha Finding 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.
gotcha The `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.

Install

pip install opt-einsum PyPI

Imports

contract
```
from opt_einsum import contract
```
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.

Quickstart

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}")

view raw JSON →