-
Notifications
You must be signed in to change notification settings - Fork 561
[Fori_loop|While_loop] Placeholder lower torch.while_loop with python dispatch for simple addition test case #6532
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
[Fori_loop|While_loop] Placeholder lower torch.while_loop with python dispatch for simple addition test case #6532
Conversation
f35899a to
122d6e5
Compare
|
CI failed due to |
mbzomowski
left a comment
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.
LGTM for the TPU CI portion: https://github.com/pytorch/xla/actions/runs/7980785965/job/21791172250
But I'll leave it to pt/xla folks to approve.
| from torch._higher_order_ops.while_loop import while_loop_op | ||
|
|
||
|
|
||
| @while_loop_op.py_impl(DispatchKey.XLA) |
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.
Great to see that this is working!
yeounoh
left a comment
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.
LGTM, left some minor comments.
… dispatch for simple addition test case (pytorch#6532)
This is the step one PR:
Due to
torch.while_loopis defined in torch/_higher_order_ops/while_loop.py, not described in native_functions.yaml, lowertorch.while_loopvia python dispatch like PR#5763TODO: current status need to patch thepython dispatchin PyTorch code, we need modify this PR to enable the python dispatch from PyTorch/XLA codechild PR: #6529
PR resource: #6527