Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Suggestion: Multiple einsum steps #216

Open
davystrong opened this issue May 26, 2023 · 2 comments
Open

Suggestion: Multiple einsum steps #216

davystrong opened this issue May 26, 2023 · 2 comments

Comments

@davystrong
Copy link

davystrong commented May 26, 2023

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
Copy link
Owner

dgasmith commented Jun 1, 2023

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
Copy link
Author

davystrong commented Jun 2, 2023

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants