Skip to content

Fix: [Inductor] Enforce strict negative index check for index_select (fixes #169779)#169788

Closed
dumko2001 wants to merge 3 commits into
pytorch:mainfrom
dumko2001:fix/inductor-index-select-negative-indices
Closed

Fix: [Inductor] Enforce strict negative index check for index_select (fixes #169779)#169788
dumko2001 wants to merge 3 commits into
pytorch:mainfrom
dumko2001:fix/inductor-index-select-negative-indices

Conversation

@dumko2001

@dumko2001 dumko2001 commented Dec 7, 2025 •
edited by pytorch-bot Bot
Loading

Copy link
Copy Markdown
Contributor

Summary
Fixes #169779.

This PR fixes an inconsistency where torch.compile (Inductor) handled negative indices in index_select incorrectly (silently wrapping them or producing wrong results), whereas Eager mode strictly raises a RuntimeError.

The Fix

  • Modified torch/_refs/__init__.py: Added a strict _assert_async check for non-negative indices in the reference implementation.
  • Why this approach? aten.index_select is already defined in the Reference implementation. Adding a separate decomposition caused registry conflicts. Modifying _refs ensures that Inductor (and any other backend using Refs) adheres to PyTorch's strict non-negative index semantics.
  • (Note: Previous attempts to modify decompositions.py or lowering.py have been reverted in favor of this cleaner Reference fix.)

Testing

  • Added regression test InductorIndexSelectTests to test/inductor/test_torchinductor.py.
  • Verified locally that compiled index_select now raises RuntimeError on negative indices.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo @mlazos

@pytorch-bot

pytorch-bot Bot commented Dec 7, 2025 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/169788

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 52e4b24 with merge base 3e026f2 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@dumko2001

dumko2001 commented Dec 7, 2025

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "bug" "release notes: inductor"

@dumko2001

dumko2001 commented Dec 7, 2025

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "release notes: inductor" "topic: not user facing"

@pytorch-bot

pytorch-bot Bot commented Dec 7, 2025

Copy link
Copy Markdown

Didn't find following labels among repository labels: bug

@dumko2001 dumko2001 mentioned this pull request Dec 7, 2025
Open
malfet
malfet requested changes Dec 8, 2025
View reviewed changes

@malfet malfet left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would work, but IMO this is a wrong approach, one should just add support for exception in torch.compile for MPS

@dumko2001 dumko2001 force-pushed the fix/inductor-index-select-negative-indices branch from d70da5f to af6cb6d Compare December 9, 2025 06:22
@dumko2001

dumko2001 commented Dec 9, 2025

Copy link
Copy Markdown
Contributor Author

@malfet Thank you for the feedback. I agree that a custom Inductor lowering is the wrong approach and have updated the PR to rely on decomposition:

  1. Reverted the custom lowering from torch/_inductor/lowering.py.
  2. Modified the core aten.index_select decomposition in torch/_decomp/decompositions.py to explicitly enforce the non-negative index constraint.

The fix now ensures the strict semantics of index_select are captured universally at the graph level (preventing the unintended wrapping of aten.index).

For the check itself, I used torch.ops.aten._assert_async.msg (as torch._check failed during tracing).

I agree that improving systemic exception propagation for torch.compile is important, but believe this decomposition fix is the necessary first step for this operator.

Please let me know if this updated approach looks better for merging.

@dumko2001

dumko2001 commented Dec 18, 2025

Copy link
Copy Markdown
Contributor Author

@malfet I have updated the PR based on the findings.

  1. I pivoted to implementing the check in torch/_refs/__init__.py. I found that index_select is already defined there, so adding a separate decomposition caused registry conflicts.
  2. I reverted the changes to decompositions.py (which caused the CI timeout in the previous run).
  3. I added the regression test inside the proper TestCase class in test_torchinductor.py.

@albanD albanD added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Dec 18, 2025
@dumko2001 dumko2001 force-pushed the fix/inductor-index-select-negative-indices branch from 3cadd4a to 52e4b24 Compare December 22, 2025 10:02
@dumko2001 dumko2001 requested a review from malfet December 22, 2025 10:03
@github-actions

github-actions Bot commented Feb 20, 2026

Copy link
Copy Markdown
Contributor

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

@github-actions github-actions Bot added the Stale label Feb 20, 2026
@github-actions github-actions Bot closed this Mar 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@malfet malfet Awaiting requested review from malfet

Assignees

No one assigned

Labels

module: inductor open source release notes: inductor Stale topic: not user facing topic category triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Inductor] Inconsistent handling of negative indices in index_select: CPU/CUDA/MPS

4 participants

@dumko2001 @malfet @albanD @pytorchbot