Improve debug_print formatting

pytensor/printing.py

@@ -291,7 +291,7 @@ N.B.:
     for var in inputs_to_print:
         _debugprint(
             var,
-            prefix="-",
+            prefix="→ ",
             depth=depth,
             done=done,
             print_type=print_type,
@@ -342,11 +342,17 @@ N.B.:
 
     if len(inner_graph_vars) > 0:
         print("", file=_file)
-        new_prefix = " >"
-        new_prefix_child = " >"
+        prefix = ""
+        new_prefix = prefix + " ← "
+        new_prefix_child = prefix + "   "
         print("Inner graphs:", file=_file)
 
+        printed_inner_graphs_nodes = set()
         for ig_var in inner_graph_vars:
+            if ig_var.owner in printed_inner_graphs_nodes:
+                continue
+            else:
+                printed_inner_graphs_nodes.add(ig_var.owner)
             # This is a work-around to maintain backward compatibility
             # (e.g. to only print inner graphs that have been compiled through
             # a call to `Op.prepare_node`)
@@ -385,6 +391,7 @@ N.B.:
 
             _debugprint(
                 ig_var,
+                prefix=prefix,
                 depth=depth,
                 done=done,
                 print_type=print_type,
@@ -399,13 +406,14 @@ N.B.:
                 print_op_info=print_op_info,
                 print_destroy_map=print_destroy_map,
                 print_view_map=print_view_map,
+                is_inner_graph_header=True,
             )
 
             if print_fgraph_inputs:
                 for inp in inner_inputs:
                     _debugprint(
                         inp,
-                        prefix="-",
+                        prefix=" → ",
                         depth=depth,
                         done=done,
                         print_type=print_type,
@@ -485,6 +493,7 @@ def _debugprint(
     parent_node: Optional[Apply] = None,
     print_op_info: bool = False,
     inner_graph_node: Optional[Apply] = None,
+    is_inner_graph_header: bool = False,
 ) -> TextIO:
     r"""Print the graph represented by `var`.
 
@@ -625,7 +634,10 @@ def _debugprint(
         else:
             data = ""
 
-        var_output = f"{prefix}{node.op}{output_idx}{id_str}{type_str}{var_name}{destroy_map_str}{view_map_str}{o}{data}"
+        if is_inner_graph_header:
+            var_output = f"{prefix}{node.op}{id_str}{destroy_map_str}{view_map_str}{o}"
+        else:
+            var_output = f"{prefix}{node.op}{output_idx}{id_str}{type_str}{var_name}{destroy_map_str}{view_map_str}{o}{data}"
 
         if print_op_info and node not in op_information:
             op_information.update(op_debug_information(node.op, node))
@@ -633,7 +645,7 @@ def _debugprint(
         node_info = (
             parent_node and op_information.get(parent_node)
         ) or op_information.get(node)
-        if node_info and var in node_info:
+        if node_info and var in node_info and not is_inner_graph_header:
             var_output = f"{var_output} ({node_info[var]})"
 
         if profile and profile.apply_time and node in profile.apply_time:
@@ -660,12 +672,13 @@ def _debugprint(
         if not already_done and (
             not stop_on_name or not (hasattr(var, "name") and var.name is not None)
         ):
-            new_prefix = prefix_child + " |"
-            new_prefix_child = prefix_child + " |"
+            new_prefix = prefix_child + " ├─ "
+            new_prefix_child = prefix_child + " │ "
 
             for in_idx, in_var in enumerate(node.inputs):
                 if in_idx == len(node.inputs) - 1:
-                    new_prefix_child = prefix_child + "  "
+                    new_prefix = prefix_child + " └─ "
+                    new_prefix_child = prefix_child + "   "
 
                 if hasattr(in_var, "owner") and hasattr(in_var.owner, "op"):
                     if (
@@ -698,6 +711,8 @@ def _debugprint(
                     print_view_map=print_view_map,
                     inner_graph_node=inner_graph_node,
                 )
+        elif not is_inner_graph_header:
+            print(prefix_child + " └─ ···", file=file)
     else:
         id_str = get_id_str(var)
 

tests/compile/test_builders.py

@@ -572,18 +572,18 @@ def test_debugprint():
     lines = output_str.split("\n")
 
     exp_res = """OpFromGraph{inline=False} [id A]
- |x [id B]
- |y [id C]
- |z [id D]
+ ├─ x [id B]
+ ├─ y [id C]
+ └─ z [id D]
 
 Inner graphs:
 
 OpFromGraph{inline=False} [id A]
- >Elemwise{add,no_inplace} [id E]
- > |*0-<TensorType(float64, (?, ?))> [id F]
- > |Elemwise{mul,no_inplace} [id G]
- >   |*1-<TensorType(float64, (?, ?))> [id H]
- >   |*2-<TensorType(float64, (?, ?))> [id I]
+ ← Elemwise{add,no_inplace} [id E]
+    ├─ *0-<TensorType(float64, (?, ?))> [id F]
+    └─ Elemwise{mul,no_inplace} [id G]
+       ├─ *1-<TensorType(float64, (?, ?))> [id H]
+       └─ *2-<TensorType(float64, (?, ?))> [id I]
 """
 
     for truth, out in zip(exp_res.split("\n"), lines):

tests/test_printing.py

@@ -145,12 +145,12 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace} [id 0]
-         |Elemwise{add,no_inplace} [id 1] 'C'
-         | |A [id 2]
-         | |B [id 3]
-         |Elemwise{add,no_inplace} [id 4]
-           |D [id 5]
-           |E [id 6]
+         ├─ Elemwise{add,no_inplace} [id 1] 'C'
+         │  ├─ A [id 2]
+         │  └─ B [id 3]
+         └─ Elemwise{add,no_inplace} [id 4]
+            ├─ D [id 5]
+            └─ E [id 6]
         """
     ).lstrip()
 
@@ -163,12 +163,12 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace} [id A]
-         |Elemwise{add,no_inplace} [id B] 'C'
-         | |A [id C]
-         | |B [id D]
-         |Elemwise{add,no_inplace} [id E]
-           |D [id F]
-           |E [id G]
+         ├─ Elemwise{add,no_inplace} [id B] 'C'
+         │  ├─ A [id C]
+         │  └─ B [id D]
+         └─ Elemwise{add,no_inplace} [id E]
+            ├─ D [id F]
+            └─ E [id G]
         """
     ).lstrip()
 
@@ -181,10 +181,11 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace} [id A]
-         |Elemwise{add,no_inplace} [id B] 'C'
-         |Elemwise{add,no_inplace} [id C]
-           |D [id D]
-           |E [id E]
+         ├─ Elemwise{add,no_inplace} [id B] 'C'
+         │  └─ ···
+         └─ Elemwise{add,no_inplace} [id C]
+            ├─ D [id D]
+            └─ E [id E]
         """
     ).lstrip()
 
@@ -196,12 +197,12 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace}
-         |Elemwise{add,no_inplace} 'C'
-         | |A
-         | |B
-         |Elemwise{add,no_inplace}
-           |D
-           |E
+         ├─ Elemwise{add,no_inplace} 'C'
+         │  ├─ A
+         │  └─ B
+         └─ Elemwise{add,no_inplace}
+            ├─ D
+            └─ E
         """
     ).lstrip()
 
@@ -213,10 +214,10 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace} 0 [None]
-         |A [None]
-         |B [None]
-         |D [None]
-         |E [None]
+         ├─ A [None]
+         ├─ B [None]
+         ├─ D [None]
+         └─ E [None]
         """
     ).lstrip()
 
@@ -231,10 +232,10 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace} 0 [None]
-         |A [None]
-         |B [None]
-         |D [None]
-         |E [None]
+         ├─ A [None]
+         ├─ B [None]
+         ├─ D [None]
+         └─ E [None]
         """
     ).lstrip()
 
@@ -249,10 +250,10 @@ def test_debugprint():
     reference = dedent(
         r"""
         Elemwise{add,no_inplace} 0 [None]
-         |A [None]
-         |B [None]
-         |D [None]
-         |E [None]
+         ├─ A [None]
+         ├─ B [None]
+         ├─ D [None]
+         └─ E [None]
         """
     ).lstrip()
 
@@ -274,26 +275,26 @@ def test_debugprint():
     exp_res = dedent(
         r"""
         Elemwise{Composite} 4
-         |InplaceDimShuffle{x,0} v={0: [0]} 3
-         | |CGemv{inplace} d={0: [0]} 2
-         |   |AllocEmpty{dtype='float64'} 1
-         |   | |Shape_i{0} 0
-         |   |   |B
-         |   |TensorConstant{1.0}
-         |   |B
-         |   |<TensorType(float64, (?,))>
-         |   |TensorConstant{0.0}
-         |D
-         |A
+         ├─ InplaceDimShuffle{x,0} v={0: [0]} 3
+         │  └─ CGemv{inplace} d={0: [0]} 2
+         │     ├─ AllocEmpty{dtype='float64'} 1
+         │     │  └─ Shape_i{0} 0
+         │     │     └─ B
+         │     ├─ TensorConstant{1.0}
+         │     ├─ B
+         │     ├─ <TensorType(float64, (?,))>
+         │     └─ TensorConstant{0.0}
+         ├─ D
+         └─ A
 
         Inner graphs:
 
         Elemwise{Composite}
-         >add
-         > |<float64>
-         > |sub
-         >   |<float64>
-         >   |<float64>
+         ← add
+            ├─ <float64>
+            └─ sub
+               ├─ <float64>
+               └─ <float64>
         """
     ).lstrip()
 
@@ -314,10 +315,10 @@ def test_debugprint_id_type():
     s = s.getvalue()
 
     exp_res = f"""Elemwise{{add,no_inplace}} [id {e_at.auto_name}]
- |dot [id {d_at.auto_name}]
- | |<TensorType(float64, (?, ?))> [id {b_at.auto_name}]
- | |<TensorType(float64, (?,))> [id {a_at.auto_name}]
- |<TensorType(float64, (?,))> [id {a_at.auto_name}]
+ ├─ dot [id {d_at.auto_name}]
+ │  ├─ <TensorType(float64, (?, ?))> [id {b_at.auto_name}]
+ │  └─ <TensorType(float64, (?,))> [id {a_at.auto_name}]
+ └─ <TensorType(float64, (?,))> [id {a_at.auto_name}]
     """
 
     assert [l.strip() for l in s.split("\n")] == [
@@ -351,15 +352,15 @@ def test_debugprint_inner_graph():
     lines = output_str.split("\n")
 
     exp_res = """MyInnerGraphOp [id A]
- |3 [id B]
- |4 [id C]
+ ├─ 3 [id B]
+ └─ 4 [id C]
 
 Inner graphs:
 
 MyInnerGraphOp [id A]
- >op2 [id D] 'igo1'
- > |*0-<MyType()> [id E]
- > |*1-<MyType()> [id F]
+ ← op2 [id D] 'igo1'
+    ├─ *0-<MyType()> [id E]
+    └─ *1-<MyType()> [id F]
     """
 
     for exp_line, res_line in zip(exp_res.split("\n"), lines):
@@ -375,19 +376,19 @@ MyInnerGraphOp [id A]
     lines = output_str.split("\n")
 
     exp_res = """MyInnerGraphOp [id A]
- |5 [id B]
+ └─ 5 [id B]
 
 Inner graphs:
 
 MyInnerGraphOp [id A]
- >MyInnerGraphOp [id C]
- > |*0-<MyType()> [id D]
- > |*1-<MyType()> [id E]
+ ← MyInnerGraphOp [id C]
+    ├─ *0-<MyType()> [id D]
+    └─ *1-<MyType()> [id E]
 
 MyInnerGraphOp [id C]
- >op2 [id F] 'igo1'
- > |*0-<MyType()> [id D]
- > |*1-<MyType()> [id E]
+ ← op2 [id F] 'igo1'
+    ├─ *0-<MyType()> [id D]
+    └─ *1-<MyType()> [id E]
     """
 
     for exp_line, res_line in zip(exp_res.split("\n"), lines):