Extend torch function support to ALL arguments, not just scalar type (but not insides of list) #145089

ezyang · 2025-01-17T17:28:38Z

Stack from ghstack (oldest at bottom):

-> Extend torch function support to ALL arguments, not just scalar type (but not insides of list) #145089

Signed-off-by: Edward Z. Yang [email protected]

cc @albanD

[ghstack-poisoned]

pytorch-bot · 2025-01-17T17:28:43Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit ac58da9 with merge base 2507ae6 ():

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, linux.12xlarge, unstable) (gh) (#158876)
/var/lib/jenkins/workspace/xla/torch_xla/csrc/runtime/BUILD:476:14: Compiling torch_xla/csrc/runtime/xla_util_test.cpp failed: (Exit 1): gcc failed: error executing CppCompile command (from target //torch_xla/csrc/runtime:xla_util_test) /usr/bin/gcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections ... (remaining 229 arguments skipped)

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

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 389c821 Pull Request resolved: #145089

ezyang · 2025-01-24T15:54:31Z

The specific problem relates to torch function redispatch. Consider the argument overload list for Tensor.set_:

  static PythonArgParser parser(
      {
          "set_()",
          "set_(Storage source)",
          "set_(Storage source, SymInt storage_offset, SymIntArrayRef size, SymIntArrayRef stride=None)",
          "set_(Tensor source)",
          "set_(Tensor source, SymInt storage_offset, SymIntArrayRef size, SymIntArrayRef stride=None)",
      },
      /*traceable=*/false);

This says we attempt to resolve the Storage overload before the Tensor overload. When a Tensor has a __torch_function__ on it, the current implementation of the patch causes it to match the Storage (because an object with torch function always matches). Now the problem is when we redispatch with torch function disabled, it STILL matches and we end up attempting the Storage overload.

So maybe we need some special case saying that if something is a Tensor, it doesn't match for non-Tensor arguments? Hmm...

[ghstack-poisoned]

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 17fc269 Pull Request resolved: #145089

[ghstack-poisoned]

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: dcef210 Pull Request resolved: #145089

ezyang · 2025-01-25T03:49:34Z

@vmoens @albanD I don't really know if I want to land this before I go on leave, it might have unintended side effects. Also, it doesn't work with entries in lists, which probably should be done before this is considered feature complete.

vmoens · 2025-01-25T18:22:11Z

Sure I'm in no hurry for this to land, we can talk about the implications once you're back!

ezyang · 2025-08-07T14:45:32Z

We ballin' this

[ghstack-poisoned]

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: de8ef4a Pull Request resolved: #145089

albanD

yolo I guess? Not sure if we want more extensive testing for this..

albanD · 2025-08-07T15:58:28Z

torch/csrc/utils/python_arg_parser.cpp

+  // that internally
+  if (check_has_torch_function(obj, /*ignore_mode*/ true) &&
+      !THPVariable_Check(obj)) {
+    // tensor subclasses and unrelated objects with __torch_function__


This needs updating as only non-subclass are caught here actually.

[ghstack-poisoned]

ezyang · 2025-08-07T20:20:52Z

@pytorchbot merge

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 3046eb6 Pull Request resolved: #145089

pytorchmergebot · 2025-08-07T20:23:29Z

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

…(but not insides of list) (pytorch#145089) Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch#145089 Approved by: https://github.com/albanD, https://github.com/zou3519

Update

e79bd4f

[ghstack-poisoned]

ezyang mentioned this pull request Jan 17, 2025

Add support for torch function on dtype arguments #145085

Closed

github-actions bot requested review from SherlockNoMad, albanD, antoniojkim, bdhirsh and miladm January 17, 2025 17:28

ezyang requested a review from vmoens January 21, 2025 14:36

Update

bc9552c

[ghstack-poisoned]

pytorch-bot bot added the release notes: fx release notes category label Jan 24, 2025

Update

de1b4fe

[ghstack-poisoned]

ezyang changed the title ~~[POC] Extend torch function support to ALL arguments, not just scalar type (but not insides of list)~~ Extend torch function support to ALL arguments, not just scalar type (but not insides of list) Jan 25, 2025

ezyang changed the title ~~Extend torch function support to ALL arguments, not just scalar type (but not insides of list)~~ [POC] Extend torch function support to ALL arguments, not just scalar type (but not insides of list) Jan 25, 2025

ezyang added the no-stale label Jan 25, 2025

ezyang changed the title ~~[POC] Extend torch function support to ALL arguments, not just scalar type (but not insides of list)~~ Extend torch function support to ALL arguments, not just scalar type (but not insides of list) Aug 7, 2025

Update

6a45ce6

[ghstack-poisoned]

ezyang added a commit that referenced this pull request Aug 7, 2025

[POC] Extend torch function support to ALL arguments, not just scalar…

13000b4

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: de8ef4a Pull Request resolved: #145089

ezyang added topic: new features topic category module: python frontend For issues relating to PyTorch's Python frontend and removed no-stale labels Aug 7, 2025

ezyang requested review from zou3519 and removed request for vmoens August 7, 2025 14:55

albanD approved these changes Aug 7, 2025

View reviewed changes

zou3519 approved these changes Aug 7, 2025

View reviewed changes

ezyang mentioned this pull request Aug 7, 2025

Delete Python reference implementation from torchdim, as it is untested #160115

Closed

Update

ac58da9

[ghstack-poisoned]

ezyang added a commit that referenced this pull request Aug 7, 2025

[POC] Extend torch function support to ALL arguments, not just scalar…

95c9929

… type (but not insides of list) Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 3046eb6 Pull Request resolved: #145089

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 7, 2025

pytorchmergebot added the merging label Aug 7, 2025

pytorchmergebot added the Merged label Aug 7, 2025

pytorchmergebot closed this in 03b254e Aug 7, 2025

pytorchmergebot removed the merging label Aug 7, 2025

ezyang mentioned this pull request Aug 10, 2025

torch_function objects passed as non-Tensor args should trigger overrides #119194

Open

github-actions bot deleted the gh/ezyang/3068/head branch September 7, 2025 02:14

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Extend torch function support to ALL arguments, not just scalar type (but not insides of list) #145089

Extend torch function support to ALL arguments, not just scalar type (but not insides of list) #145089

Uh oh!

ezyang commented Jan 17, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jan 17, 2025 •

edited

Loading

Uh oh!

ezyang commented Jan 24, 2025

Uh oh!

ezyang commented Jan 25, 2025

Uh oh!

vmoens commented Jan 25, 2025

Uh oh!

ezyang commented Aug 7, 2025

Uh oh!

albanD left a comment

Uh oh!

albanD Aug 7, 2025

Uh oh!

ezyang commented Aug 7, 2025

Uh oh!

pytorchmergebot commented Aug 7, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Extend torch function support to ALL arguments, not just scalar type (but not insides of list) #145089

Extend torch function support to ALL arguments, not just scalar type (but not insides of list) #145089

Uh oh!

Conversation

ezyang commented Jan 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jan 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

ezyang commented Jan 24, 2025

Uh oh!

ezyang commented Jan 25, 2025

Uh oh!

vmoens commented Jan 25, 2025

Uh oh!

ezyang commented Aug 7, 2025

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Aug 7, 2025

Choose a reason for hiding this comment

Uh oh!

ezyang commented Aug 7, 2025

Uh oh!

pytorchmergebot commented Aug 7, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ezyang commented Jan 17, 2025 •

edited

Loading

pytorch-bot bot commented Jan 17, 2025 •

edited

Loading