Repository navigation
Fix einsum ellipsis lowering to MIL - #2706
Merged
TobyRoseman merged 2 commits intoMay 22, 2026
Merged
Conversation
parse_einsum_equation maps each character to a unique int, so every "." in "..." becomes index 0 and solve_diagonal_einsum reads them as three duplicate indices, producing a perm longer than the input rank. Expand "..." to fresh single-character labels in build_einsum_mil before parsing, using each input var's rank to count how many dims its ellipsis covers. Fixes apple#2644
TobyRoseman
requested changes
May 22, 2026
TobyRoseman
left a comment
TobyRoseman
left a comment
Collaborator
There was a problem hiding this comment.
Overall, this change looks good. Two minor comments.
Collaborator
TobyRoseman
approved these changes
May 22, 2026
Collaborator
|
Thanks for your contribution @eyupcanakman |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
torch.einsum("...,f->...f", x, y)fails to convert withperm should have the same length as rank(x): 3 != 2. Any equation with...hits the same path.parse_einsum_equationmaps each character to a unique int. Every.in...becomes index 0, sosolve_diagonal_einsumreads them as three duplicate indices and produces apermlonger than the input rank.The fix replaces
...with fresh single-character labels before that call, using each input var's rank to count how many dims it covers. The replacement happens inbuild_einsum_mil, which already receives the input vars. Equations without...are unaffected.test_einsum_with_ellipsiscovers the reported...,f->...fshape and fails on main without the fix.Fixes #2644