Add section about .code to docs

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/18493

Differential Revision: D14634677

Pulled By: jamesr66a

fbshipit-source-id: 9ee065f6ce4218f725b93deb4c64b4ef55926145

diff --git a/docs/source/jit.rst b/docs/source/jit.rst
index 6b2880f..ad968a0 100644 (file)
--- a/docs/source/jit.rst
+++ b/docs/source/jit.rst
@@ -694,8 +694,56 @@ Disable JIT for Debugging
     function.
 
 
+Inspecting Code
+^^^^^^^^^^^^^^^
+
+    TorchScript provides a code pretty-printer for all ScriptModule instances. This
+    pretty-printer gives an interpretation of the script method's code as valid
+    Python syntax. For example::
+
+        @torch.jit.script
+        def foo(len):
+            # type: (int) -> torch.Tensor
+            rv = torch.zeros(3, 4)
+            for i in range(len):
+                if i < 10:
+                    rv = rv - 1.0
+                else:
+                    rv = rv + 1.0
+                return rv
+
+        print(foo.code)
+
+    A ``ScriptModule`` with a single ``forward`` method will have an attribute
+    ``code``, which you can use to inspect the ``ScriptModule``'s code.
+    If the ScriptModule has more than one method, you will need to access
+    ``.code`` on the method itself and not the module. We can inspect the
+    code of a method named ``bar`` on a ScriptModule by accessing ``.bar.code``.
+
+    The example script abouve produces the code::
+
+        def forward(self,
+                    len: int) -> Tensor:
+            rv = torch.zeros([3, 4], dtype=None, layout=None, device=None)
+            rv0 = rv
+            for i in range(len):
+                if torch.lt(i, 10):
+                    rv1 = torch.sub(rv0, 1., 1)
+                else:
+                    rv1 = torch.add(rv0, 1., 1)
+                rv0 = rv1
+            return rv0
+
+    This is TorchScript's interpretation of the code for the ``forward`` method.
+    You can use this to ensure TorchScript (tracing or scripting) has captured
+    your model code correctly.
+
+
 Interpreting Graphs
 ^^^^^^^^^^^^^^^^^^^
+    TorchScript also has a representation at a lower level than the code pretty-
+    printer, in the form of IR graphs.
+
     TorchScript uses a static single assignment (SSA) intermediate representation
     (IR) to represent computation. The instructions in this format consist of
     ATen (the C++ backend of PyTorch) operators and other primitive operators,
@@ -714,11 +762,8 @@ Interpreting Graphs
 
         print(foo.graph)
 
-    A ``ScriptModule`` with a single ``forward`` method will have an attribute
-    ``graph``, which you can use to inspect the IR representing the computation.
-    If the ScriptModule has more than one method, you will need to access
-    ``.graph`` on the method itself and not the module. We can inspect the
-    graph of a method named ``bar`` on a ScriptModule by accessing ``.bar.graph``.
+    ``.graph`` follows the same rules described in the Inspecting Code section
+    with regard to ``forward`` method lookup.
 
     The example script above produces the graph::
 

diff --git a/torch/csrc/jit/script/init.cpp b/torch/csrc/jit/script/init.cpp
index 5ecebf3..0692763 100644 (file)
--- a/torch/csrc/jit/script/init.cpp
+++ b/torch/csrc/jit/script/init.cpp
@@ -987,7 +987,16 @@ void initJitScriptBindings(PyObject* module) {
         std::vector<ClassTypePtr> classes;
         PythonPrint(oss, m, constants, classes, true);
         return std::make_pair(oss.str(), std::move(constants));
-      });
+      })
+      .def_property_readonly(
+          "code",
+          [](Method& self) {
+            std::ostringstream ss;
+            std::vector<at::Tensor> tensors;
+            std::vector<ClassTypePtr> classes;
+            PythonPrint(ss, self, tensors, classes, false);
+            return ss.str();
+          });
 
   m.def(
       "_jit_script_compile",