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
Fix sync on index_select decomposition when the index has one element #125973
Conversation
The previous fix was not general enough. Fixes #125952 [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125973
Note: Links to docs will display an error until the docs builds have been completed. ❌ 24 New FailuresAs of commit 001ca95 with merge base 2ed17e0 (): NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
@@ -4020,10 +4020,10 @@ def index_select(x: TensorLike, dim: int, index: TensorLike): | |||
) | |||
if index.ndim == 0: | |||
index = index.unsqueeze(0) | |||
if x.ndim == 0: | |||
if x.numel() == 1: | |||
# Treat scalars as elements of \R^1 | |||
# We cannot use x[idx] here as it accesses item() (??), hence this awkward construction |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we know why this happens?!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
See #105641
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we add an operator that is like x[idx]
but instead produces a copy rather than view?
This seems reasonable but I guess it doesn't work? |
The roll error is fine, but that pdist error I'll have to look into. I'll spin up an A100 box next week and make sure the fix I submit works lol |
And it's also difficult to tell how this could fix it, the trace appears to originate from C++ code |
I thought the issue was coming from the use of index_select and its decomposition dispatching to advanced indexing, which sometime synchronizes |
Stack from ghstack (oldest at bottom):
The previous fix was not general enough.
Fixes #125952