Adds input/output to the text rendering (#109)

* Adds input/output to the text rendering

* mypy

Author: xadupre

_unittests/ut_helpers/test_graph_helper.py

@@ -126,21 +126,31 @@ def test_text_rendering(self):
         )
         graph = GraphRendering(proto)
         text = textwrap.dedent(graph.text_rendering(prefix="|")).strip("\n")
+        print()
+        print(text)
         expected = textwrap.dedent(
             """
             |
             |
             |
             |
-            |   Add             Neg
-            |    |               |
-            |    +-------+-------+
-            |            |
-            |           Mul
-            |            |
-            |            +-------+
-            |                    |
-            |                   Mul
+            |   X            Y
+            |   |            |
+            |   +------+-----+-----------+
+            |          |                 |
+            |         Add               Neg
+            |          |                 |
+            |          +-----+-----------+
+            |                |
+            |               Mul
+            |                |
+            |                +-----------+
+            |                            |
+            |                           Mul
+            |                            |
+            |                            +----------------------+
+            |                                                   |
+            |                                                   Z
             |
             |
             """

_unittests/ut_torch_models/test_tiny_llms_onnx.py

@@ -28,14 +28,15 @@ def test_onnx_export_tiny_llm_official(self):
         self.assertEqual(
             {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
         )
-        ep = torch.onnx.export(
-            model,
-            (),
-            kwargs=inputs,
-            dynamic_shapes=data["dynamic_shapes"],
-            dynamo=True,
-            optimize=True,
-        )
+        with torch_export_patches(patch_transformers=True):
+            ep = torch.onnx.export(
+                model,
+                (),
+                kwargs=inputs,
+                dynamic_shapes=data["dynamic_shapes"],
+                dynamo=True,
+                optimize=True,
+            )
         # There are some discrepancies with torch==2.6
         if not has_torch("2.7"):
             raise unittest.SkipTest("discrepancies observed with torch<2.7")

onnx_diagnostic/helpers/graph_helper.py

@@ -1,5 +1,6 @@
 from typing import Dict, List, Optional, Sequence, Set, Tuple, Union
 import onnx
+import onnx.helper as oh
 
 
 class GraphRendering:
@@ -146,6 +147,34 @@ def start_names(self) -> List[onnx.NodeProto]:
         )
         return [*input_names, *init_names]
 
+    @property
+    def input_names(self) -> List[str]:
+        "Returns the list of input names."
+        return (
+            self.proto.input
+            if isinstance(self.proto, onnx.FunctionProto)
+            else [
+                i.name
+                for i in (
+                    self.proto if isinstance(self.proto, onnx.GraphProto) else self.proto.graph
+                ).input
+            ]
+        )
+
+    @property
+    def output_names(self) -> List[str]:
+        "Returns the list of output names."
+        return (
+            self.proto.output
+            if isinstance(self.proto, onnx.FunctionProto)
+            else [
+                i.name
+                for i in (
+                    self.proto if isinstance(self.proto, onnx.GraphProto) else self.proto.graph
+                ).output
+            ]
+        )
+
     @classmethod
     def build_node_edges(cls, nodes: Sequence[onnx.NodeProto]) -> Set[Tuple[int, int]]:
         """Builds the list of edges between nodes."""
@@ -210,7 +239,7 @@ def text_edge(
         else:
             middle = (p1[1] + p2[1]) // 2
             a, b = (p1[1] + 1, middle - 1) if p1[1] < middle else (middle + 1, p1[1] - 1)
-            grid[p1[0] + 2][a : b + 1] = ["-" * (b - a + 1)]
+            grid[p1[0] + 2][a : b + 1] = ["-"] * (b - a + 1)
             a, b = (p1[1] + 1, middle - 1) if p1[1] < middle else (middle + 1, p1[1] - 1)
             grid[p1[0] + 2][a : b + 1] = ["-"] * (b - a + 1)
 
@@ -223,7 +252,6 @@ def text_edge(
     def text_rendering(self, prefix="") -> str:
         """
         Renders a model in text.
-        It only renders nodes.
 
         .. runpython::
             :showcode:
@@ -257,8 +285,14 @@ def text_rendering(self, prefix="") -> str:
             text = textwrap.dedent(graph.text_rendering()).strip("\\n")
             print(text)
         """
-        nodes = self.nodes
-        existing = self.start_names
+        nodes = [
+            *[oh.make_node(i, ["BEGIN"], [i]) for i in self.input_names],
+            *self.nodes,
+            *[oh.make_node(i, [i], ["END"]) for i in self.output_names],
+        ]
+        exist = set(self.start_names) - set(self.input_names)
+        exist |= {"BEGIN"}
+        existing = sorted(exist)
         order = self.computation_order(nodes, existing)
         positions = self.graph_positions(nodes, order, existing)
         text_pos = self.text_positions(nodes, positions)
@@ -269,8 +303,10 @@ def text_rendering(self, prefix="") -> str:
 
         for n1, n2 in edges:
             self.text_edge(grid, text_pos[n1], text_pos[n2])
+            assert len(set(len(g) for g in grid)) == 1, f"lengths={[len(g) for g in grid]}"
         for node, pos in zip(nodes, text_pos):
             self.text_grid(grid, pos, node.op_type)
+            assert len(set(len(g) for g in grid)) == 1, f"lengths={[len(g) for g in grid]}"
 
         return "\n".join(
             f"{prefix}{line.rstrip()}" for line in ["".join(line) for line in grid]