csarofeen
diff --git a/‎test/test_dynamic_shapes.py
Lines changed: 32 additions & 0 deletions b/‎test/test_dynamic_shapes.py
Lines changed: 32 additions & 0 deletions
diff --git a/‎torch/fx/graph.py
Lines changed: 37 additions & 3 deletions b/‎torch/fx/graph.py
Lines changed: 37 additions & 3 deletions
@@ -9,7 +9,9 @@
 import torch
 import operator
 import itertools
+import io
 from torch.utils._pytree import tree_map
+from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx.experimental.symbolic_shapes import ShapeEnv, PySymInt, sym_float
 from torch.utils._python_dispatch import TorchDispatchMode
 
@@ -354,5 +356,35 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
 
         self.assertTrue(sym_int_encountered)
 
+    @skipIfNoSympy
+    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
+    def test_print_readable_with_symints(self, mock_stdout):
+        def f(a, b):
+            dim0 = a.shape[0] + b.shape[0]
+            dim1 = a.shape[1] + b.shape[1]
+            d = a.new_empty(dim0, dim1)
+            d = torch.ops.aten.native_dropout(d, 0.5, train=True)
+            return d
+
+        fx_g = make_fx(f, tracing_mode="symbolic")(torch.randn(5, 3), torch.randn(4, 3))
+        fx_g.print_readable()
+
+        self.assertExpectedInline(mock_stdout.getvalue().strip(), """\
+class f(torch.nn.Module):
+    def forward(self, a_1: f32[t0.size(0),t0.size(1)], b_1: f32[t1.size(0),t0.size(1)]):
+        # No stacktrace found for following nodes
+        sym_size: Sym(t0.size(0)) = torch.ops.aten.sym_size(a_1, 0)
+        sym_size_1: Sym(t1.size(0)) = torch.ops.aten.sym_size(b_1, 0)
+        add: Sym(t0.size(0) + t1.size(0)) = sym_size + sym_size_1;  sym_size = sym_size_1 = None
+        sym_size_2: Sym(t0.size(1)) = torch.ops.aten.sym_size(a_1, 1)
+        sym_size_3: Sym(t0.size(1)) = torch.ops.aten.sym_size(b_1, 1);  b_1 = None
+        add_1: Sym(2*t0.size(1)) = sym_size_2 + sym_size_3;  sym_size_2 = sym_size_3 = None
+        new_empty: f32[t0.size(0) + t1.size(0),2*t0.size(1)] = torch.ops.aten.new_empty.default(a_1, [add, add_1], dtype = torch.float32, layout = torch.strided, device = device(type='cpu'), pin_memory = False);  a_1 = add = add_1 = None
+        native_dropout = torch.ops.aten.native_dropout.default(new_empty, 0.5, True);  new_empty = None
+        getitem: f32[t0.size(0) + t1.size(0),2*t0.size(1)] = native_dropout[0]
+        getitem_1: b8[t0.size(0) + t1.size(0),2*t0.size(1)] = native_dropout[1];  native_dropout = None
+        return (getitem, getitem_1)""")  # noqa: B950
+
+
 if __name__ == '__main__':
     run_tests()
@@ -2,8 +2,8 @@
 import torch.utils._pytree as pytree
 from . import _pytree as fx_pytree
 from ._compatibility import compatibility
-import contextlib
 
+import contextlib
 from typing import TYPE_CHECKING, Callable, Any, List, Dict, NamedTuple, Optional, Tuple, Set, FrozenSet, Type
 from dataclasses import dataclass
 from contextlib import contextmanager
@@ -188,6 +188,21 @@ def _is_illegal_name(self, name: str, obj: Any) -> bool:
 
         return False
 
+dtype_abbrs = {
+    torch.bfloat16: 'bf16',
+    torch.float64: 'f64',
+    torch.float32: 'f32',
+    torch.float16: 'f16',
+    torch.complex32: 'c32',
+    torch.complex64: 'c64',
+    torch.complex128: 'c128',
+    torch.int8: 'i8',
+    torch.int16: 'i16',
+    torch.int32: 'i32',
+    torch.int64: 'i64',
+    torch.bool: 'b8',
+    torch.uint8: 'u8',
+}
 
 @compatibility(is_backward_compatible=True)
 @dataclass
@@ -457,10 +472,29 @@ def append_stacktrace_summary(node : Node):
                         body.append(f'\n# {summary_str}\n')
                 elif prev_stacktrace != "":
                     prev_stacktrace = ""
-                    body.append('\n# No stacktrace found for following nodes \n')
+                    body.append('\n# No stacktrace found for following nodes\n')
+
+        def stringify_shape(shape : torch.Size) -> str:
+            return f"[{','.join(str(x) for x in shape)}]"
 
         def emit_node(node : Node):
             maybe_type_annotation = '' if node.type is None else f' : {type_repr(node.type)}'
+
+            if verbose:
+                # override annotation with more detailed information
+                from torch._subclasses.fake_tensor import FakeTensor
+                from torch.fx.experimental.proxy_tensor import _py_sym_types
+                from torch.fx.passes.shape_prop import TensorMetadata
+
+                meta_val = node.meta.get('val', node.meta.get('tensor_meta', None))
+
+                if isinstance(meta_val, FakeTensor):
+                    maybe_type_annotation = f': {dtype_abbrs[meta_val.dtype]}{stringify_shape(meta_val.shape)}'
+                elif isinstance(meta_val, _py_sym_types):
+                    maybe_type_annotation = f': Sym({meta_val.expr})'
+                elif isinstance(meta_val, TensorMetadata):
+                    maybe_type_annotation = f': {dtype_abbrs[meta_val.dtype]}{stringify_shape(meta_val.shape)}'
+
             if node.op == 'placeholder':
                 assert isinstance(node.target, str)
                 maybe_default_arg = '' if not node.args else f' = {repr(node.args[0])}'
@@ -552,7 +586,7 @@ def emit_node(node : Node):
 
         prologue = self.gen_fn_def(free_vars, maybe_return_annotation[0])
 
         code = ''.join(body)
+        code = ''.join(body).lstrip('\n')
         code = '\n'.join('    ' + line for line in code.split('\n'))
         fn_code = f"""
 {wrap_stmts}