dgasmith / opt_einsum

⚡️Optimizing einsum functions in NumPy, Tensorflow, Dask, and more with contraction order optimization.
https://dgasmith.github.io/opt_einsum/
MIT License
822 stars 67 forks source link

Suggestion: Multiple einsum steps #216

Open davystrong opened 1 year ago

davystrong commented 1 year ago

Hello, I've been working on a little project to compile multiple einsum subscripts into a single operation, and was wondering if there would be any interest in integrating it into opt_einsum.

Maybe this is something highly specific that most people don't need, but the basic idea is that if you have multiple np.einsum steps they can frequently be combined into a single step, thus avoiding intermediate arrays which are normally undesirable. This can frequently be done even if the output is reshaped between steps (depending on the shapes) by carefully reshaping the input array.

This should be fully backwards compatible with Numpy's einsum implementation and I thought of seeing if there was interest in adding it to that, but it feels like a better fit here (also, I switched to using opt_einsum internally as I was so impressed with the performance jump). Obviously it would require significant rework around the parsing steps to work, but I'd be interested to know if there's any interest. But maybe this is just be a really specific problem that most people don't face, and should be kept in its own package.

dgasmith commented 1 year ago

This is very interesting and a request that we have had previously. I would be open to merging the functionality here if we have sufficient tests and clear documentation for it.

What do you think @jcmgray?

davystrong commented 1 year ago

In terms of tests, the good news is that getting the expected output is very simple since the naive implementation is just repeatedly running einsum. However actually insuring test coverage is much more difficult, although I have tried to add any errors as test cases during development. I'll try to think of a more systematic set of tests. Though any implementation which modifies the existing contract function would require a switch for a single argument, so anything I miss should only affect new code, not the existing contract code. Still, better not have any bugs.

One other question: in my current implementation I include a fallback when shapes are incompatible which runs einsum more than once reshaping the output in between. I try to split the subscripts at the point where the intermediate arrays are smallest. Is this worth adding in this package too?