[scan] scan dim handling in user-facing scan() #145179
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/145179
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 93c417e with merge base d53f206 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
torch/_higher_order_ops/scan.py
Outdated
| assert to_dim >= 0 and to_dim < t.ndim | ||
| assert from_dim >= 0 and from_dim < t.ndim |
There was a problem hiding this comment.
| assert to_dim >= 0 and to_dim < t.ndim | |
| assert from_dim >= 0 and from_dim < t.ndim | |
| assert 0 <= to_dim < t.ndim | |
| assert 0 <= from_dim < t.ndim |
Or seperate it into two seperate assertion statements to make it easier to figure out which constraint is violated
There was a problem hiding this comment.
Thank you for your suggestion. I have refactored the code as you suggested.
|
@pytorchbot label "topic: not user facing" |
torch/_higher_order_ops/scan.py
Outdated
| # https://github.com/pytorch/pytorch/pull/139864 is merged | ||
| # say we have a tensor of shape [3, 4, 5, 6] | ||
| # shift_source_dim_to_target_dim(t, 0, 3) -> [4, 5, 6, 3] | ||
| def shift_source_dim_to_target_dim(t, from_dim: int, to_dim: int): |
There was a problem hiding this comment.
We probably won't need the wrapper. Can we use torch.movedim directly (for both PRs)?
There was a problem hiding this comment.
Thank you for the catch, I fixed that in that PR as well as in the one for the associative_scan
|
@ydwu4, thank you for your review. Could you please trigger the CI tests for this PR and also maybe look at the associative_scan PR? |
|
@ydwu4 I did a minor update, maybe could you trigger the CI tests? |
|
@pytorchbot merge |
Merge startedYour 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 |
|
The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command |
Merge startedYour 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 |
Merge failedReason: Comment with id 2625390323 not found Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge |
Merge startedYour 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 |
This PR introduces the capability that the scan dim is handled in the user facing scan() call. Internally, the scan dim is always shifted to dim 0 and then the scan is performed over that dim.
This is a follow-up PR from bohnstingl#3
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @ydwu4