Need help with opt_einsum?
Click the “chat” button below for chat support from the developer who created it, or find similar developers for support.

About the developer

dgasmith
412 Stars 35 Forks MIT License 330 Commits 25 Opened issues

Description

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

Services available

!
?

Need anything else?

Contributors list

Optimized Einsum

Build Status codecov Anaconda-Server Badge PyPI PyPIStats Documentation Status DOI

Optimized Einsum: A tensor contraction order optimizer

Optimized einsum can significantly reduce the overall execution time of einsum-like expressions (e.g.,

np.einsum
,
dask.array.einsum
,
pytorch.einsum
,
tensorflow.einsum
, ) by optimizing the expression's contraction order and dispatching many operations to canonical BLAS, cuBLAS, or other specialized routines.

Optimized einsum is agnostic to the backend and can handle NumPy, Dask, PyTorch, Tensorflow, CuPy, Sparse, Theano, JAX, and Autograd arrays as well as potentially any library which conforms to a standard API. See the documentation for more information.

Example usage

The

opt_einsum.contract
function can often act as a drop-in replacement for

einsum
functions without further changes to the code while providing superior performance. Here, a tensor contraction is preformed with and without optimization:
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

In this particular example, we see a ~3000x performance improvement which is not uncommon when compared against unoptimized contractions. See the backend examples for more information on using other backends.

Features

The algorithms found in this repository often power the

einsum
optimizations in many of the above projects. For example, the optimization of
np.einsum
has been passed upstream and most of the same features that can be found in this repository can be enabled with
np.einsum(..., optimize=True)
. However, this repository often has more up to date algorithms for complex contractions.

The following capabilities are enabled by

opt_einsum
:

Please see the documentation for more features!

Installation

opt_einsum
can either be installed via
pip install opt_einsum
or from conda
conda install opt_einsum -c conda-forge
. See the installation documentation for further methods.

Citation

If this code has benefited your research, please support us by citing:

Daniel G. A. Smith and Johnnie Gray, opt_einsum - A Python package for optimizing contraction order for einsum-like expressions. Journal of Open Source Software, 2018, 3(26), 753

DOI: https://doi.org/10.21105/joss.00753

Contributing

All contributions, bug reports, bug fixes, documentation improvements, enhancements, and ideas are welcome.

A detailed overview on how to contribute can be found in the contributing guide.

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.