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

Support using matmul for Einsum equations with 3 or more inputs #300

Merged
merged 1 commit into from
Aug 7, 2024

Conversation

robertknight
Copy link
Owner

Previously Einsum equations with more than two inputs were handled via a fallback to a naive implementation that was much less efficient than a matmul if a reduction was required. In this commit, equations are split into a path where each step handles one or two inputs. Each step will use matmul, multiply or reduce-sum as appropriate.

The generated path always processes inputs in left-to-right order. This can be much less efficient than the optimal contraction order. Using a better algorithm is left as a future exercise.

Part of #298.

@robertknight robertknight force-pushed the einsum-faster-many-input branch 2 times, most recently from 8b6e0c9 to 3c5c3b8 Compare August 7, 2024 05:40
@robertknight robertknight mentioned this pull request Aug 7, 2024
7 tasks
@robertknight robertknight marked this pull request as ready for review August 7, 2024 05:46
Previously Einsum equations with more than two inputs were handled via a
fallback to a naive implementation that was much less efficient than a matmul if
a reduction was required. In this commit, equations are split into a path where
each step handles one or two inputs. Each step will use matmul, multiply or
reduce-sum as appropriate.

The generated path always processes inputs in left-to-right order. This can be
much less efficient than the optimal contraction order. Using a better algorithm
is left as a future exercise.
@robertknight robertknight merged commit cc9abc3 into main Aug 7, 2024
2 checks passed
@robertknight robertknight deleted the einsum-faster-many-input branch August 7, 2024 05:55
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

Successfully merging this pull request may close these issues.

1 participant