opt-einsum-fx
Einsum optimization using opt_einsum and PyTorch FX
Decision gist · record as of 2026-08-14
No. The package is abandoned (last commit 2022-03-17, no releases since 2021-11-07) and likely incompatible with current PyTorch and opt_einsum versions. While the idea is sound and the MIT license is permissive, the lack of maintenance means you would need to fork and patch it yourself to use it with modern libraries. For new projects, consider using opt_einsum directly or PyTorch's built-in optimization passes instead.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires PyTorch with FX support and opt_einsum installed; symbolic_trace requires the function to be traceable (no dynamic control flow).
- Low install friction with a pure-Python wheel.
- However, the package is abandoned—last commit was 2022-03-17 and no releases since 2021-11-07.
License · maintenance · safety
MIT (permissive) — MIT license permits commercial and private use with minimal restrictions, requiring only license and copyright notice preservation.
last release 2021-11-07 (1741 days) · last repo commit 2022-03-17 · 22 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 504,342 downloads/mo, #6,298 on PyPI
Alternatives
Verify before relying
import torch
import torch.fx
import opt_einsum_fx
def einmatvecmul(a, b, vec):
return torch.einsum("zij,zjk,zk->zi", a, b, vec)
graph_mod = torch.fx.symbolic_trace(einmatvecmul)
graph_opt = opt_einsum_fx.optimize_einsums_full(
model=graph_mod,
example_inputs=(torch.randn(7, 4, 5), torch.randn(7, 5, 3), torch.randn(7, 3))
)- Compatibility with PyTorch versions released after 2022; FX API may have evolved.
- Whether the package works with current opt_einsum releases without modification.
- Performance gains on modern hardware (GPU, recent CPU architectures).
What it is and what it does
opt_einsum_fx wraps PyTorch's FX graph tracing to automatically optimize einsum operations within traced functions. It uses opt_einsum to reorder tensor contractions for better computational efficiency, then rewrites the traced graph with the optimized contraction sequence. The package is designed for developers working with batched tensor operations, matrix products, and other multi-tensor contractions expressed via einsum notation.
The typical workflow is to define a function containing einsum calls, trace it with torch.fx.symbolic_trace, pass the traced module to optimize_einsums_full with example inputs, and receive an optimized module with reordered contractions. The fact sheet shows a concrete example achieving roughly 2x speedup on CPU for a batched matrix-vector product, though gains depend on tensor shapes and hardware.
Use it for
- Accelerate batched matrix operations in neural networks by reordering einsum contraction sequences.
- Optimize custom tensor algebra kernels expressed via einsum without rewriting the contraction logic.
- Profile and improve performance of multi-tensor operations in research code before deployment.
- Automatically find better contraction orders for complex tensor expressions in scientific computing.
- Reduce computational cost of attention mechanisms or other einsum-heavy layers in transformer models.
Worth the install?
AI-flagged interpretation of the facts on this page. Verify before relying on it.
No.
The package is abandoned (last commit 2022-03-17, no releases since 2021-11-07) and likely incompatible with current PyTorch and opt_einsum versions. While the idea is sound and the MIT license is permissive, the lack of maintenance means you would need to fork and patch it yourself to use it with modern libraries. For new projects, consider using opt_einsum directly or PyTorch's built-in optimization passes instead.
Install
opt-einsum-fx on PyPI
Before you install
Low install friction with a pure-Python wheel. However, the package is abandoned—last commit was 2022-03-17 and no releases since 2021-11-07. It may not be compatible with recent PyTorch or opt_einsum versions without manual updates.
Requires PyTorch with FX support and opt_einsum installed; symbolic_trace requires the function to be traceable (no dynamic control flow).
License in practice
MIT license permits commercial and private use with minimal restrictions, requiring only license and copyright notice preservation.
Quickstart
import torch
import torch.fx
import opt_einsum_fx
def einmatvecmul(a, b, vec):
return torch.einsum("zij,zjk,zk->zi", a, b, vec)
graph_mod = torch.fx.symbolic_trace(einmatvecmul)
graph_opt = opt_einsum_fx.optimize_einsums_full(
model=graph_mod,
example_inputs=(torch.randn(7, 4, 5), torch.randn(7, 5, 3), torch.randn(7, 3))
)Verify before relying
- Compatibility with PyTorch versions released after 2022; FX API may have evolved.
- Whether the package works with current opt_einsum releases without modification.
- Performance gains on modern hardware (GPU, recent CPU architectures).
Package facts
| License | MIT permissive |
| Python support | Supports the current Python release >=3.6 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 3 packagestorchopt-einsumpackaging |
| Maintenance | Abandoned 1,741 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 504,342 / month, #6,298 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaLicense :: OSI Approved :: MIT LicenseOperating System :: OS IndependentProgramming Language :: Python :: 3 |
Evidence: opt_einsum_fx-0.1.4-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “einsum optimization pytorch”
- opt-einsum-fxOptimizes PyTorch einsum operations and functions containing them by…
- opt-einsumOptimizes the contraction order of einsum expressions to dramatically…
- cotengraCotengra optimizes the contraction order of large tensor networks and…
Give your agent the search over MCP, or paste the wish link into any chat.
More Scientific/Engineering packages
NumPy provides an N-dimensional array object and a comprehensive suite of mathematical, linear algebra, Fourier transform, and random number functions for scientific computing in Python.
pandas provides fast, flexible data structures (Series and DataFrame) for loading, cleaning, transforming, and analyzing labeled or relational data in Python.
scipy provides numerical algorithms for mathematics, science, and engineering—including optimization, integration, linear algebra, Fourier transforms, signal and image processing, and ODE solvers—built on numpy arrays.
scikit-learn provides a comprehensive Python library for supervised and unsupervised machine learning, including classification, regression, clustering, dimensionality reduction, and model evaluation tools built on NumPy and SciPy.
Install it if you need to train, evaluate, or deploy supervised or unsupervised learning models.
dill extends Python's pickle module to serialize and deserialize a much wider range of Python objects, including functions, lambdas, classes, and interpreter sessions, to byte streams for storage or network transmission.
Multiprocess is an enhanced fork of Python's standard multiprocessing library that uses dill for better serialization, allowing you to spawn processes with a threading-like API and share complex objects between them.
Install it if you use multiprocessing and encounter pickle serialization limits with lambdas or complex objects.
See also opt-einsum · cotengra · spo4onnx · depyf · einops · thop · torch-einops-utils · lion-pytorch · powerfx · codeflash