Library / SDK
dgasmith/opt_einsum avatar
dgasmith/opt_einsum

opt_einsum picks the contraction order so your tensor code stops guessing

⚡️Optimizing einsum functions in NumPy, Tensorflow, Dask, and more with contraction order optimization.

992 stars82 forksPythonMIT

At a glance

What is it?
opt_einsum is a path optimizer for einsum style expressions: it works out which order to contract tensors in, then dispatches the individual steps to BLAS, cuBLAS or whatever backend you plug in. It is a drop-in for np.einsum in the simplest case, and a deeper tool when you want to inspect the path, reuse it, or share intermediates across many contractions.
Who is it for?
opt_einsum is a good fit for scientific Python code with heavy tensor contractions, especially where the same expression is evaluated repeatedly and the path can be computed once and reused, and it is a good fit for teams that need the chosen path printed out for review.
Can I use it commercially?
Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
Is it still maintained?
Yes. The repository last received commits 10 days ago.
What is it written in?
Mainly Python, according to GitHub's language statistics.

Answers come from the project's GitHub data, last synced on November 7, 2026, and from our analysis. They are not legal advice.

From daily snapshots

Star history

  • dgasmith/opt_einsum992 ★ · +1
Star history
Repository2026-10-212026-11-07
dgasmith/opt_einsum991992

1 Stars in the last 17 days, measured from the daily GitHub snapshot. Compare with other projects

Editorial analysis

The problem is the order of the contractions, not the arithmetic

Every array library will multiply and sum your tensors if you ask it to, but the order in which those operations are carried out changes the cost enormously. Contracting a chain of matrices left to right and right to left can differ by orders of magnitude in work, and the best order depends on the shapes involved. opt_einsum is a library for finding that order. It reads the subscripts of an expression, builds a graph of possible pairwise contractions, searches for a good sequence, and then dispatches each step to a canonical BLAS, cuBLAS or other specialised routine. The project is agnostic to the backend by design and can handle NumPy, Dask, PyTorch, TensorFlow, CuPy, Sparse, Theano, JAX and Autograd arrays, and the documentation suggests that any library conforming to a standard API could work too. The package itself is MIT licensed, written in Python, and requires Python 3.12 or newer according to its project metadata.

The README benchmark, and how to read it honestly

The documentation gives one worked comparison, and it is worth reading with care because the numbers are the authors own and come from one machine and one expression. The example builds a set of random arrays, contracts five of them with a single einsum call, and times it against the optimized version:

python
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)

%timeit np.einsum('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
1 loops, best of 3: 934 ms per loop

%timeit contract('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
1000 loops, best of 3: 324 us per loop

The documentation calls this a performance improvement of roughly 3000 times and notes that gains of that size are not unusual when an unoptimized contraction is compared with an optimized one. What the example does not tell you is how the numbers change with shape, with the number of terms in the expression, or with a backend other than NumPy, so treat it as an illustration of the gap rather than as a benchmark you can quote for your own workload. The `contract` function is the main entry point and, in the simple case, is meant to be usable as a drop-in replacement for an existing einsum call with no other code changes.

Installing it takes one command from pip, uv or conda-forge

The installation section is short because the package has no runtime dependencies at all, which the project metadata confirms with an empty dependency list. The default route is pip:

bash
pip install opt_einsum

The same package is also published on conda-forge, and the documentation shows the uv route as well:

bash
uv pip install opt_einsum

Backends arrive as extras, which matters because the package ships without a tensor library of its own. The optional dependency groups cover dask, jax, numpy with a floor of 1.23, sparse, and torch with a range that stops below 3.0, and separate groups exist for testing, linting, documentation and a combined dev extra. For development, the documentation suggests cloning the repository and installing in editable mode with the dev extra:

bash
uv pip install -e ".[dev]"

The repository ships a Makefile that drives the same workflow through uv, with a format target that runs ruff check and ruff format, a mypy target, a test target that runs pytest with coverage, and a docs target that builds the MkDocs site. So the loop for a contributor is uv, ruff, mypy, pytest and MkDocs, all pinned in the project configuration.

Backend support is a dispatch decision, not a rewrite

The interesting architectural choice is that opt_einsum does not implement the contractions itself. It decides what to contract next, and the array library does the work. That is why the same optimizer can drive a CPU BLAS call in NumPy, a GPU routine in CuPy, or a distributed array in Dask, and why the package can be used with frameworks it has never heard of as long as they follow a standard API. The cost of this design is that an expression which is efficient on one backend may be a poor fit for another, since a good path for a CPU matrix multiplication is not automatically a good path for a distributed one. The documentation keeps backend examples on a separate page for that reason. A related nuance: the algorithms in this repository are the ones that were passed upstream into NumPy, where most of the same features are reachable through `np.einsum(..., optimize=True)`, which is the route to take if you want the standard library behaviour and no extra dependency, and the reason to install opt_einsum directly is that the project usually carries more up to date algorithms for complex contractions.

Paths can be inspected, reused and compiled

Beyond speed, the library exposes the decisions it makes, which matters for anything that runs in production. The documentation describes inspecting detailed information about the path that was chosen, so a reviewer can see why a particular sequence was picked rather than trusting it. The second capability is reuse: expressions, potentially with constant tensors attached, can be compiled for greater performance, and the reuse documentation covers specifying those constants. A compiled path can then be applied to many calls with the same shapes, which is the pattern that removes the search cost from the hot loop entirely. The third capability is scale: an arbitrary number of indices can be used to find contractions for expressions involving hundreds or even thousands of tensors, and the examples section points at a greedy path finder for exactly that case. So the library is not only a drop-in speedup for a single call, it is a tool for turning a one-off expression into a compiled, inspectable unit you can test.

Shared intermediates and gradients close the loop

Two more features decide whether opt_einsum fits a larger system. The first is the ability to share intermediate computations among multiple contractions, so when several expressions reuse the same sub-contraction the work is done once instead of once per expression. That is the difference between speeding up one call and speeding up a model, and it interacts with the compiled path idea, because a path that knows about the shared piece can plan around it. The second is differentiation: the library can compute gradients of tensor contractions using autograd or JAX. That means an optimized contraction can sit inside a differentiable pipeline without the optimization breaking the gradient, which is a real consideration for machine learning code where a contraction is often only one step in a larger graph. The practical caveat is that both features only pay off when the workload is repetitive and the shapes are stable, since the path search itself is not free and is only worth amortising if you can cache its result.

Maintenance, packaging and the citation to keep

The project is stable rather than experimental, and the packaging reflects that: the project metadata carries the Development Status classifier for Production/Stable, supports CPython and PyPy, and declares support for Python 3.12, 3.13 and 3.14. The build uses hatchling with a VCS derived version written into `opt_einsum/_version.py`, and the version and readme are declared dynamic, so the tag is the single source of truth for the release number. On maintenance, the pattern is worth noticing: the last push to the repository was on 2026-10-15, while the most recent tagged release listed is v3.4.0 from 2024-09-26, with v3.3.0 from 2020-07-19 and v3.2.1 from 2020-04-15 before it. In other words the code keeps moving without a stream of tagged releases, so if you depend on a specific behaviour, pin a version rather than tracking the default branch. The project asks that research users cite the paper by Daniel G. A. Smith and Johnnie Gray in the Journal of Open Source Software, 2018, volume 3, issue 26, article 753.

Editorial conclusion

opt_einsum is a good fit for scientific Python code with heavy tensor contractions, especially where the same expression is evaluated repeatedly and the path can be computed once and reused, and it is a good fit for teams that need the chosen path printed out for review. It is not a fit for small contractions, where the search cost can exceed the work it saves, and it is not a substitute for understanding the algebra behind the expression, because the optimizer picks an order, not a formula. Before you adopt it, check three things: whether your backend is one of the supported ones or conforms to the standard array API, whether the path can be cached for the shapes you actually run, and the version you pin, since the most recent tagged release is v3.4.0 from 2024-09-26 while the last push to the repository was on 2026-10-15.

Frequently asked questions

How do I install opt_einsum?

With `pip install opt_einsum`, with `uv pip install opt_einsum`, or from conda-forge with `conda install opt_einsum -c conda-forge`. The package itself has no runtime dependencies, and backends are installed as extras such as numpy, torch, jax, dask and sparse.

Can I use opt_einsum as a drop-in replacement for np.einsum?

In the simple case, yes. The `contract` function takes the same subscripts and operands and is described as often acting as a drop-in replacement for einsum functions without further code changes, while still allowing you to inspect the path it chose.

Does opt_einsum work with PyTorch and TensorFlow?

Yes. The library is backend agnostic and supports NumPy, Dask, PyTorch, TensorFlow, CuPy, Sparse, Theano, JAX and Autograd arrays, and it can also compute gradients of tensor contractions using autograd or JAX.

When should I not use opt_einsum?

The path search is not free, so very small contractions are unlikely to benefit and the overhead can outweigh the saving. If NumPy already gives you what you need, most of the same features are reachable through `np.einsum(..., optimize=True)`, and the project itself notes it usually carries more up to date algorithms for complex contractions.

What is the latest version of opt_einsum?

The most recent release listed is v3.4.0, published on 2024-09-26, and the last push to the repository was on 2026-10-15. The version number is derived from the VCS tag by the build configuration, so pinning a tag is the way to fix a version.

Official sources

  1. dgasmith/opt_einsum on GitHub
  2. License: MIT
  3. Project website
  4. README
  5. Releases
Add this badge to your README

If you maintain this project, the badge below links readers to this analysis and shows its maintenance status from the daily GitHub snapshot. Paste the markdown into your README; add ?metric=license or ?metric=stars to the image URL for a different field.

Add this badge to your README

markdown
[![Hysen Labs](https://hysenlabs.com/badge/dgasmith-opt-einsum.svg)](https://hysenlabs.com/projects/dgasmith-opt-einsum)

Star history chart for your README

markdown
[![Star history](https://hysenlabs.com/chart/dgasmith-opt-einsum.svg)](https://hysenlabs.com/projects/dgasmith-opt-einsum)

Was this useful?