Improvements for associative_scan - Reverse feature #133011
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/133011
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1e3391a with merge base 161cc13 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| leaves, spec = pytree.tree_flatten(input) | ||
|
|
||
| if reverse: | ||
| leaves = [torch.flip(elem, [dim]) for elem in leaves] |
There was a problem hiding this comment.
In practice elem.flip(dim) also works (no need to construct a list [dim]) but unfortunately this is not documented properly:
ydwu4
left a comment
ydwu4
left a comment
There was a problem hiding this comment.
Overall looks good to me. Left some comments for test organizations. We generally don't want unit test to be too complicated or take too long to run and we should split a large test into different tests that test different things.
test/functorch/test_control_flow.py
Outdated
| print("Flip test fails for backends: " + str(fails_for_backend)) | ||
| self.assertEqual(len(fails_for_backend), 0) | ||
|
|
||
| for n in range(20): |
There was a problem hiding this comment.
20 might be too large? Don't want the unit tests take too long
There was a problem hiding this comment.
Well, the problem is that with the combine_fn='generic' that will follow, the test will randomly fail at some point. Thus I wanted to keep it like this. However, if you feel that it takes too long, I can reduce.
There was a problem hiding this comment.
Do we know why it fails randomly?
There was a problem hiding this comment.
This test fails for n=9 for the case combine_mode='generic' . The output of
torch.compile(
associative_scan, backend=backend, fullgraph=True
)
operation is
tensor([36, 36, 35, 33, 30, 26, 21, 15, 8], device='cuda:0') which is correct, while the output of
associative_scan2 = associative_scan
is
tensor([36, 15, 35, 0, 30, 0, 21, 0, 8], device='cuda:0'), which is incorrect.
This problem only appears for n=9, while for other values of n, the outputs are correct.
|
Can skip the failed test with |
zou3519
added
the
triaged
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see 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 |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
@pytorchbot merge -i |
Merge startedYour change will be merged while ignoring the following 0 checks: Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
This is part of a series of PRs to improve the functionality of the `associatve_scan` functionality. This specific PR introduces a `reverse` flag to the `associative_scan` to establish a similar interface as for `jax.associative_scan`. This PR has been derived from #129307. @ydwu4 @Chillee @zou3519 Pull Request resolved: #133011 Approved by: https://github.com/ydwu4
6 participants
This is part of a series of PRs to improve the functionality of the
associatve_scanfunctionality. This specific PR introduces areverseflag to theassociative_scanto establish a similar interface as forjax.associative_scan. This PR has been derived from #129307.@ydwu4 @Chillee @zou3519
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang