From 938bc1062fe26999a2ac765361eefdb29fb156f4 Mon Sep 17 00:00:00 2001
From: Ansley Ussery <ansley@fb.com>
Date: Thu, 21 Jan 2021 12:22:00 -0800
Subject: [PATCH] Add alternative prettyprinting method to `Graph`

ghstack-source-id: 06b05e5ae284ce026a6aa39c0933706d03a54f2a
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50878
---
 torch/fx/examples/inspect_utils.py | 14 --------------
 torch/fx/graph.py                  | 16 ++++++++++++++++
 2 files changed, 16 insertions(+), 14 deletions(-)
 delete mode 100644 torch/fx/examples/inspect_utils.py

diff --git a/torch/fx/examples/inspect_utils.py b/torch/fx/examples/inspect_utils.py
deleted file mode 100644
index 26833c1c6e14..000000000000
--- a/torch/fx/examples/inspect_utils.py
+++ /dev/null
@@ -1,14 +0,0 @@
-from tabulate import tabulate
-
-"""
-The methods in this file may be used to examine the state of the code
-and how the Graph evolves at any time during execution. If you're
-unsure of what's happening in an example in this folder, try adding one
-of these methods before and after a key line.
-"""
-
-def print_IR(graph):
-    node_specs = [[n.op, n.name, n.target, n.args, n.kwargs]
-                  for n in graph.nodes]
-    print(tabulate(node_specs,
-                   headers=['opcode', 'name', 'target', 'args', 'kwargs']))
diff --git a/torch/fx/graph.py b/torch/fx/graph.py
index 09e7f2bfc8d9..45b084011962 100644
--- a/torch/fx/graph.py
+++ b/torch/fx/graph.py
@@ -785,6 +785,22 @@ def format_node(n : Node) -> Optional[str]:
                 s += '\n    ' + node_str
         return s
 
+    def print_tabular(self):
+        """
+        Prints the intermediate representation of the graph in tabular
+        format.
+        """
+        try:
+            from tabulate import tabulate
+        except ImportError:
+            print("`print_tabular` relies on the library `tabulate`, "
+                  "which could not be found on this machine. Run `pip "
+                  "install tabulate` to install the library.")
+        node_specs = [[n.op, n.name, n.target, n.args, n.kwargs]
+                      for n in self.nodes]
+        print(tabulate(node_specs,
+              headers=['opcode', 'name', 'target', 'args', 'kwargs']))
+
     def lint(self, root : Optional[torch.nn.Module] = None):
         """
         Runs various checks on this Graph to make sure it is well-formed. In