Support list output for HigherOrderOperators #101986
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/101986
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 32d8904: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
ydwu4
added
ciflow/trunk
torch/_dynamo/variables/torch.py
Outdated
| return isinstance(var, TensorVariable) or ( | ||
| isinstance(var, ConstantVariable) and var.value is None | ||
| ) | ||
|
|
||
| if any( | ||
| not _is_supported(var) for var in output.unpack_var_sequence(tx) | ||
| ): | ||
| unimplemented( | ||
| "HigherOrderOperator body's output must consist of tensors or None only" | ||
| ) |
There was a problem hiding this comment.
Do we really need to support None output?
There was a problem hiding this comment.
yeah, the reason I added is because the backward logic of autograd.Function can consist of None and we don't want it to graph break. The specific example I encountered is this: https://github.com/pytorch/pytorch/blob/main/test/dynamo/test_repros.py#L2918, where XSoftmax's backward returns None.
There was a problem hiding this comment.
I'm confused, because when we trace the backward, always_restore should be True, so we should never get here
There was a problem hiding this comment.
Proposal: the logic should be something like:
if always_restore:
# we don't care about what the output is
pass
else:
# we require the output to be a list/tuple of Tensors
# If it isn't then we raise notimplemented
When always_restore is True, we don't actually care what the output is, because we are never going to use it to construct a subgraph. I don't want to allow non-Tensor output in higher order ops just yet but it might be inevitable.
|
@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 |
Fixes the issue in #100278: support list output for HigherOrderOperator.
cc @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx