Add expanded_def option for FX printing, render descriptor, update tests #158708

ezyang · 2025-07-19T16:40:18Z

Stack from ghstack (oldest at bottom):

First, we add a new expanded_def to FX, which will expand the
definitions of variables into multiple lines, one per variable
definition. This makes extremely long args/return lists much
more readable.
Next, we extend this mechanism to also print out descriptors on
placeholders and return values, as comments, if available. This
is how we will test descriptors.
We update tlparse for AOTAutograd to use this format.
We update expect tests to use this format and update their formats,
so you can inspect what it can look at. There may be other tests
I should update, open to suggestions.

Signed-off-by: Edward Z. Yang ezyang@meta.com

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

[ghstack-poisoned]

- First, we add a new expanded_def to FX, which will expand the definitions of variables into multiple lines, one per variable definition. This makes extremely long args/return lists much more readable. - Next, we extend this mechanism to also print out descriptors on placeholders and return values, as comments, if available. This is how we will test descriptors. - We update tlparse for AOTAutograd to use this format. - We update expect tests to use this format and update their formats, so you can inspect what it can look at. There may be other tests I should update, open to suggestions. Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 09486be Pull-Request: #158708

pytorch-bot · 2025-07-19T16:40:22Z

🔗 Helpful Links

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

📄 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 5b49756 with merge base 85ee2fb ():

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)
sccache: error: couldn't connect to server

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

[ghstack-poisoned]

test/dynamo/test_subclasses.py

ezyang · 2025-07-21T14:42:51Z

test/dynamo/test_subclasses.py

+        primals_5: "Sym(s47)",  # SubclassSizeAOTInput(base=PlainAOTInput(idx=2), idx=0)
+        primals_6: "Sym(s16)",  # SubclassSizeAOTInput(base=PlainAOTInput(idx=2), idx=1)
+        primals_7: "Sym(s16)",  # SubclassStrideAOTInput(base=PlainAOTInput(idx=2), idx=0)
+    ):
        mul: "f32[s47, s16]" = torch.ops.aten.mul.Tensor(primals_3, primals_1);  primals_3 = None
        mul_3: "f32[s47, s16]" = torch.ops.aten.mul.Tensor(primals_4, primals_1);  primals_4 = None
        return (mul, mul_3, primals_5, primals_7, primals_7, primals_1, primals_5, primals_7)


Return printing not working here either 🤔

It's because these are post partition and we lost the output desc

torch/_functorch/_aot_autograd/graph_capture.py

wconstab · 2025-07-23T21:01:16Z

torch/_functorch/_aot_autograd/graph_capture.py

+            # Unfortunately, flat_args_descs is not guaranteed to match the
+            # number of actual arguments that show up on the FX graph.
+            # Speciifcally, allow_token_discovery=True means that we will
+            # silently add extra token arguments to the backwards graph.


where can i learn what "token arguments" are?

cc @angelayi but I believe https://docs.google.com/document/d/179QyhicGzTXJ5jvTAoAosP_Nzgf3PpgZwU_E3VV9PlM/edit?tab=t.0#heading=h.pj3msnr95n30 was the design doc. But the idea is simple: if you need to prevent an operator from getting DCE'ed or reordered around other operators, having a "token" tensor which flows from one operation to the next enforces this dependency without having to introduce a new form of control dependency.

wconstab · 2025-07-23T21:02:57Z

torch/_functorch/_aot_autograd/utils.py

@@ -538,3 +538,12 @@ def call_and_expect_output_descs(fn, args):
        outs_descs,
    )
    return outs_pair
+
+
+def fn_wrappers(fn):


oh, is this the helper that you referred to as fn_stack in the next PR?

yes, and whoops I misremembered the name

wconstab

assuming you fix the returns that aren't printing correctly, this LGTM.

[ghstack-poisoned]

- First, we add a new expanded_def to FX, which will expand the definitions of variables into multiple lines, one per variable definition. This makes extremely long args/return lists much more readable. - Next, we extend this mechanism to also print out descriptors on placeholders and return values, as comments, if available. This is how we will test descriptors. - We update tlparse for AOTAutograd to use this format. - We update expect tests to use this format and update their formats, so you can inspect what it can look at. There may be other tests I should update, open to suggestions. Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 4df34ed Pull-Request: #158708

linux-foundation-easycla · 2025-07-24T20:52:00Z

The committers listed above are authorized under a signed CLA.

✅ login: ezyang / name: Edward Z. Yang (0d6d5ff, 5eb250c, 5b49756, afcd4e5, 7b5fe5a, 8ac501c, e7a6e32)

ezyang · 2025-07-24T20:52:21Z

Pushed the return values fix: the problem was I needed to have partitioner preserve descs when it creates the new output nodes. This is now implemented.

[ghstack-poisoned]

- First, we add a new expanded_def to FX, which will expand the definitions of variables into multiple lines, one per variable definition. This makes extremely long args/return lists much more readable. - Next, we extend this mechanism to also print out descriptors on placeholders and return values, as comments, if available. This is how we will test descriptors. - We update tlparse for AOTAutograd to use this format. - We update expect tests to use this format and update their formats, so you can inspect what it can look at. There may be other tests I should update, open to suggestions. Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 4df34ed Pull-Request: #158708

[ghstack-poisoned]

- First, we add a new expanded_def to FX, which will expand the definitions of variables into multiple lines, one per variable definition. This makes extremely long args/return lists much more readable. - Next, we extend this mechanism to also print out descriptors on placeholders and return values, as comments, if available. This is how we will test descriptors. - We update tlparse for AOTAutograd to use this format. - We update expect tests to use this format and update their formats, so you can inspect what it can look at. There may be other tests I should update, open to suggestions. Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 29645ba Pull-Request: #158708

ezyang · 2025-07-25T13:19:00Z

@pytorchbot merge -f "all clear everything is green"

pytorchmergebot · 2025-07-25T13:21:48Z

Starting merge as part of PR stack under #158734

pytorchmergebot · 2025-07-25T13:22:06Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

Wrapping is load bearing for things that introspect argument signatures, but use of functools.wraps to do this is undesirable as this overrides the name/module of the wrapping function, which is bad for tracking down exactly what code is actually being run at runtime. simple_wraps is like wraps but it doesn't override the name information, so you still get an appropriate printout. To see the stack of all functions wrapping each other, there is now a helper fn_stack. I also make some assertions tighter in the descriptor PR. These didn't catch any bugs but I figure might as well. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #158734 Approved by: https://github.com/wconstab ghstack dependencies: #158624, #158708

…riptors (#158715) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #158715 Approved by: https://github.com/fmassa, https://github.com/wconstab, https://github.com/xmfan ghstack dependencies: #158624, #158708, #158734

…sts (#158708) ---- - First, we add a new expanded_def to FX, which will expand the definitions of variables into multiple lines, one per variable definition. This makes extremely long args/return lists much more readable. - Next, we extend this mechanism to also print out descriptors on placeholders and return values, as comments, if available. This is how we will test descriptors. - We update tlparse for AOTAutograd to use this format. - We update expect tests to use this format and update their formats, so you can inspect what it can look at. There may be other tests I should update, open to suggestions. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #158708 Approved by: https://github.com/wconstab ghstack dependencies: #158624

Wrapping is load bearing for things that introspect argument signatures, but use of functools.wraps to do this is undesirable as this overrides the name/module of the wrapping function, which is bad for tracking down exactly what code is actually being run at runtime. simple_wraps is like wraps but it doesn't override the name information, so you still get an appropriate printout. To see the stack of all functions wrapping each other, there is now a helper fn_stack. I also make some assertions tighter in the descriptor PR. These didn't catch any bugs but I figure might as well. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #158734 Approved by: https://github.com/wconstab ghstack dependencies: #158624, #158708

…riptors (#158715) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #158715 Approved by: https://github.com/fmassa, https://github.com/wconstab, https://github.com/xmfan ghstack dependencies: #158624, #158708, #158734

Update

0d6d5ff

[ghstack-poisoned]

ezyang mentioned this pull request Jul 18, 2025

Rename modules in AOTAutograd #158449

Closed

pytorch-bot bot added ciflow/inductor module: dynamo release notes: fx release notes category labels Jul 19, 2025

This was referenced Jul 18, 2025

De-abstract premature generalization with InductorWrapper #158528

Closed

Track descriptors for all inputs/outputs of AOTAutograd traced graph #158624

Closed

github-actions bot requested review from albanD, antoniojkim, bdhirsh, miladm and SherlockNoMad July 19, 2025 16:40

facebook-github-bot added the fx label Jul 19, 2025

Update

5eb250c

[ghstack-poisoned]

ezyang requested review from avikchaudhuri, tugsbayasgalan, zhxchen17, ydwu4, angelayi and Chillee as code owners July 20, 2025 02:39

ezyang mentioned this pull request Jul 20, 2025

Add aot_export_joint_with_descriptors and aot_compile_joint_with_descriptors #158715

Closed

Update

8ac501c

[ghstack-poisoned]

ezyang mentioned this pull request Jul 21, 2025

Use simple_wraps instead of functools.wraps in AOTAutograd #158734

Closed

ezyang requested review from jamesjwu and wconstab July 21, 2025 04:28

ezyang added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 21, 2025

ezyang commented Jul 21, 2025

View reviewed changes

test/dynamo/test_subclasses.py Outdated Show resolved Hide resolved

ezyang commented Jul 21, 2025

View reviewed changes

ezyang mentioned this pull request Jul 22, 2025

[WIP] Checkpoint #158859

Closed

wconstab reviewed Jul 23, 2025

View reviewed changes

torch/_functorch/_aot_autograd/graph_capture.py Outdated Show resolved Hide resolved

wconstab reviewed Jul 23, 2025

View reviewed changes

wconstab approved these changes Jul 23, 2025

View reviewed changes

Update

7b5fe5a

[ghstack-poisoned]

This was referenced Jul 24, 2025

Add aot_autograd.fx_utils #159005

Closed

Docs on export joint with descriptors #159006

Open

ezyang force-pushed the gh/ezyang/3110/head branch from a900aa4 to 7b5fe5a Compare July 24, 2025 20:54

ezyang force-pushed the gh/ezyang/3110/base branch from 932e256 to 4fdb7f9 Compare July 24, 2025 20:54

Update

e7a6e32

[ghstack-poisoned]

Update

5b49756

[ghstack-poisoned]

pytorchmergebot added the merging label Jul 25, 2025

pytorchmergebot added the Merged label Jul 25, 2025

pytorchmergebot closed this in 204eb4d Jul 25, 2025

pytorchmergebot removed the merging label Jul 25, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add expanded_def option for FX printing, render descriptor, update tests #158708

Add expanded_def option for FX printing, render descriptor, update tests #158708

Uh oh!

ezyang commented Jul 19, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 19, 2025 •

edited

Loading

Uh oh!

Uh oh!

ezyang Jul 21, 2025

Uh oh!

ezyang Jul 24, 2025

Uh oh!

Uh oh!

wconstab Jul 23, 2025

Uh oh!

ezyang Jul 24, 2025

Uh oh!

wconstab Jul 23, 2025

Uh oh!

ezyang Jul 24, 2025

Uh oh!

wconstab left a comment

Uh oh!

linux-foundation-easycla bot commented Jul 24, 2025 •

edited

Loading

Uh oh!

ezyang commented Jul 24, 2025

Uh oh!

ezyang commented Jul 25, 2025

Uh oh!

pytorchmergebot commented Jul 25, 2025

Uh oh!

pytorchmergebot commented Jul 25, 2025

Uh oh!

Uh oh!

Add expanded_def option for FX printing, render descriptor, update tests #158708

Add expanded_def option for FX printing, render descriptor, update tests #158708

Uh oh!

Conversation

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

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 19, 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/158708

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Uh oh!

ezyang Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

ezyang Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

wconstab Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

ezyang Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

wconstab Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

ezyang Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

wconstab left a comment

Choose a reason for hiding this comment

Uh oh!

linux-foundation-easycla bot commented Jul 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ezyang commented Jul 24, 2025

Uh oh!

ezyang commented Jul 25, 2025

Uh oh!

pytorchmergebot commented Jul 25, 2025

Uh oh!

pytorchmergebot commented Jul 25, 2025

Merge started

Uh oh!

Uh oh!

ezyang commented Jul 19, 2025 •

edited

Loading

pytorch-bot bot commented Jul 19, 2025 •

edited

Loading

linux-foundation-easycla bot commented Jul 24, 2025 •

edited

Loading