Update on "Only print dde partial fx graph for export"

bobrenjc93 · bobrenjc93 · commit 2a100bb1f009 · 2025-03-23T15:33:41.000-07:00
Lazos correctly pointed out this doesn't make sense for compile since
we graph break in compile. This results in tons of unwanted user log
spew. We do want this in export though since it's drastiaclly reduced
the support load for DDEs. This PR does the refactor to keep it in
export but remove it from compile

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

[ghstack-poisoned]
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -1347,7 +1347,7 @@ def run(self):
                     ).print_readable(
                         print_output=False, include_stride=True, include_device=True
                     )
-                    e.partial_fx_graph = readable_graph
+                    e.partial_fx_graph = readable_graph  # type: ignore[attr-defined]
                     raise
 
                 raise
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
@@ -1,11 +1,11 @@
 # mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import dataclasses
-import sys
 import functools
 import inspect
 import logging
 import re
+import sys
 import time
 import warnings
 from contextlib import contextmanager, nullcontext
@@ -1098,7 +1098,7 @@ def wrapper(*args, **kwargs):
                     flags=_EXPORT_FLAGS,
                 )
 
-            if hasattr(e, 'partial_fx_graph'):
+            if hasattr(e, "partial_fx_graph"):
                 print(
                     e.partial_fx_graph,
                     file=sys.stderr,
diff --git a/torch/fx/_symbolic_trace.py b/torch/fx/_symbolic_trace.py
@@ -846,7 +846,7 @@ def forward(*args, **kwargs):
                     root_module="self",
                     verbose=True,
                 ).src
-                e.partial_fx_graph = partial_fx_graph
+                e.partial_fx_graph = partial_fx_graph  # type: ignore[attr-defined]
                 raise
 
             raise

Original file line number	Diff line number	Diff line change
`@@ -1347,7 +1347,7 @@ def run(self):`
`1347`	`1347`	`).print_readable(`
`1348`	`1348`	`print_output=False, include_stride=True, include_device=True`
`1349`	`1349`	`)`
`1350`		`- e.partial_fx_graph = readable_graph`
	`1350`	`+ e.partial_fx_graph = readable_graph # type: ignore[attr-defined]`
`1351`	`1351`	`raise`
`1352`	`1352`
`1353`	`1353`	`raise`