[Numpy-discussion] Optimizing multi-tensor contractions in numpy

Charles R Harris charlesr.harris at gmail.com
Fri Jan 2 22:21:49 EST 2015


On Fri, Jan 2, 2015 at 7:45 PM, Daniel Smith <dgasmith at icloud.com> wrote:

> Hello everyone,
>
> I have been working on a chunk of code that basically sets out to provide
> a single function that can take an arbitrary einsum expression and computes
> it in the most optimal way. While np.einsum can compute arbitrary
> expressions, there are two drawbacks to using pure einsum: einsum does not
> consider building intermediate arrays for possible reductions in overall
> rank and is not currently capable of using a vendor BLAS. I have been
> working on a project that aims to solve both issues simultaneously:
>
> https://github.com/dgasmith/opt_einsum
>
> This program first builds the optimal way to contract the tensors
> together, or using my own nomenclature a “path.” This path is then iterated
> over and uses tensordot when possible and einsum for everything else. In
> test cases the worst case scenario adds a 20 microsecond overhead penalty
> and, in the best case scenario, it can reduce the overall rank of the
> tensor. The primary (if somewhat exaggerated) example is a 5-term N^8 index
> transformation that can be reduced to N^5; even when N is very small (N=10)
> there is a 2,000 fold speed increase over pure einsum or, if using
> tensordot, a 2,400 fold speed increase. This is somewhat similar to the new
> np.linalg.multi_dot function.
>
> If you are interested in this function please head over to the github repo
> and check it out. I believe the README is starting to become
> self-explanatory, but feel free to email me with any questions.
>
> This originally started because I was looking into using numpy to rapidly
> prototype quantum chemistry codes. The results of which can be found here:
> https://github.com/dgasmith/psi4numpy
>
> As such, I am very interested in implementing this into numpy. While I
> think opt_einsum is in a pretty good place, there is still quite a bit to
> do (see outstanding issues in the README). Even if this is not something
> that would fit into numpy I would still be very interested in your comments.
>
>
Sounds interesting. I wouldn't be opposed to including an optimized einsum
in numpy, there has even been some mention of using blas. Note that cblas
is used in multiarray in the current development branch, so that might be
useful. I also looked into using einsum to implement the '@' operator
coming in Python 3.5, but there was a rather large fixed overhead involved
in parsing the input string, and the multiplications were much slower than
the numpy dot operator. If those times could be reduced einsum might become
a real possibility.

Chuck
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20150102/6a82e54a/attachment.html>


More information about the NumPy-Discussion mailing list