Skip to main content

Path optimization of einsum functions.

Project description

Optimized Einsum

Tests 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 performed 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.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

opt_einsum-3.4.0.tar.gz (63.0 kB view details)

Uploaded Source

Built Distribution

opt_einsum-3.4.0-py3-none-any.whl (71.9 kB view details)

Uploaded Python 3

File details

Details for the file opt_einsum-3.4.0.tar.gz.

File metadata

  • Download URL: opt_einsum-3.4.0.tar.gz
  • Upload date:
  • Size: 63.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.0 CPython/3.11.9

File hashes

Hashes for opt_einsum-3.4.0.tar.gz
Algorithm Hash digest
SHA256 96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac
MD5 ec38160dbeecbdcccbffd1421e4c8fff
BLAKE2b-256 8cb92ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb

See more details on using hashes here.

Provenance

File details

Details for the file opt_einsum-3.4.0-py3-none-any.whl.

File metadata

  • Download URL: opt_einsum-3.4.0-py3-none-any.whl
  • Upload date:
  • Size: 71.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.0 CPython/3.11.9

File hashes

Hashes for opt_einsum-3.4.0-py3-none-any.whl
Algorithm Hash digest
SHA256 69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd
MD5 f6b4a07ed34b08eaaa833867f9ffdbce
BLAKE2b-256 23cd066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540

See more details on using hashes here.

Provenance

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page