Skip to content

[dynamo] Add handler to constant-fold elementwise_dtypes during tracing#177743

Closed
ydwu4 wants to merge 3 commits intogh/ydwu4/393/basefrom
gh/ydwu4/393/head
Closed

[dynamo] Add handler to constant-fold elementwise_dtypes during tracing#177743
ydwu4 wants to merge 3 commits intogh/ydwu4/393/basefrom
gh/ydwu4/393/head

Conversation

@ydwu4
Copy link
Copy Markdown
Contributor

@ydwu4 ydwu4 commented Mar 18, 2026

Stack from ghstack (oldest at bottom):

Fix issue 2 discovered in #177166.

elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @chauhang @amjames @Lucaskabela @jataylo @azahed98

elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

[ghstack-poisoned]
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 18, 2026

🔗 Helpful Links

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

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

✅ No Failures

As of commit 2fc8827 with merge base a345892 (image):
💚 Looks good so far! There are no failures yet. 💚

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

ydwu4 added a commit that referenced this pull request Mar 18, 2026
elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

ghstack-source-id: f4d6f22
Pull Request resolved: #177743
@pytorch-bot pytorch-bot bot added ciflow/torchtitan Run TorchTitan integration tests module: dynamo labels Mar 18, 2026
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 18, 2026

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@ydwu4 ydwu4 added the topic: not user facing topic category label Mar 18, 2026
@ydwu4 ydwu4 added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 2, 2026
…uring tracing"


Fix issue 2 discovered in #177166.

elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo azahed98

[ghstack-poisoned]
ydwu4 added a commit that referenced this pull request Apr 2, 2026
elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

ghstack-source-id: c826f82
Pull Request resolved: #177743
…uring tracing"


Fix issue 2 discovered in #177166.

elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo azahed98

[ghstack-poisoned]
ydwu4 added a commit that referenced this pull request Apr 3, 2026
elementwise_dtypes was registered as TorchInGraphFunctionVariable via
torch._higher_order_ops.out_dtype, causing dynamo to try putting it in the
FX graph. Since it returns (dtype, dtype) rather than tensors, this failed
with "torch.* op returned non-Tensor". The fix adds a handler that evaluates
elementwise_dtypes eagerly on fake tensor metadata during compilation and
returns the result as constants.

ghstack-source-id: ba09207
Pull Request resolved: #177743
@ydwu4
Copy link
Copy Markdown
Contributor Author

ydwu4 commented Apr 8, 2026

@pytorchbot merge

@pytorchmergebot
Copy link
Copy Markdown
Collaborator

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/inductor ciflow/torchtitan Run TorchTitan integration tests ciflow/trunk Trigger trunk jobs on your pull request Merged module: dynamo topic: not user facing topic category

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants