[Nested Tensor]fix sdpa backward for the special case with ragged second batch dim and constant length

[YuqingJ]

Stack from ghstack (oldest at bottom):

• -> [Nested Tensor]fix sdpa backward for the special case with ragged second batch dim and constant length #128349

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit 9fdcb0e with merge base 920ebcc ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• trunk / macos-py3-arm64-mps / test (mps, 1, 1, macos-m1-13) (gh) (similar failure)
  test_mps.py::TestMPS::test_mps_allocator_module
• trunk / macos-py3-arm64-mps / test (mps, 1, 1, macos-m1-14) (gh) (similar failure)
  test_mps.py::TestMPS::test_mps_allocator_module

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

[jbschlosser]

Looks great, thanks for the fix!

[YuqingJ]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict gh/yuqingj/1/orig returned non-zero exit code 1

Rebasing (1/1)
Auto-merging test/test_nestedtensor.py
CONFLICT (content): Merge conflict in test/test_nestedtensor.py
error: could not apply dc2c1ba22a1... [Nested Tensor]fix sdpa backward for the special case with ragged second batch dim and constant length
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
hint: Disable this message with "git config advice.mergeConflict false"
Could not apply dc2c1ba22a1... [Nested Tensor]fix sdpa backward for the special case with ragged second batch dim and constant length

Raised by https://github.com/pytorch/pytorch/actions/runs/9522329915

[YuqingJ]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased gh/yuqingj/1/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/128349)

[YuqingJ]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

• pull / linux-jammy-py3.8-gcc11 / test (distributed, 2, 2, linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[YuqingJ]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased gh/yuqingj/1/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/128349)

[YuqingJ]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here