Add print_shape and print_memory_map option to debugprint

pytensor/printing.py

@@ -89,6 +89,7 @@
     | Sequence[Variable | Apply | Function | FunctionGraph],
     depth: int = -1,
     print_type: bool = False,
+    print_shape: bool = False,
     file: Literal["str"] | TextIO | None = None,
     id_type: IDTypesType = "CHAR",
     stop_on_name: bool = False,
@@ -123,6 +124,8 @@
         Print graph to this depth (``-1`` for unlimited).
     print_type
         If ``True``, print the `Type`\s of each `Variable` in the graph.
+    print_shape
+        If ``True``, print the shape of each `Variable` in the graph.
     file
         When `file` extends `TextIO`, print to it; when `file` is
         equal to ``"str"``, return a string; when `file` is ``None``, print to
@@ -265,6 +268,7 @@
             depth=depth,
             done=done,
             print_type=print_type,
+            print_shape=print_shape,
             file=_file,
             id_type=id_type,
             inner_graph_ops=inner_graph_vars,
@@ -295,6 +299,7 @@
             depth=depth,
             done=done,
             print_type=print_type,
+            print_shape=print_shape,
             file=_file,
             topo_order=topo_order,
             id_type=id_type,
@@ -365,6 +370,7 @@
                 depth=depth,
                 done=done,
                 print_type=print_type,
+                print_shape=print_shape,
                 file=_file,
                 id_type=id_type,
                 inner_graph_ops=inner_graph_vars,
@@ -387,6 +393,7 @@
                         depth=depth,
                         done=done,
                         print_type=print_type,
+                        print_shape=print_shape,
                         file=_file,
                         id_type=id_type,
                         stop_on_name=stop_on_name,
@@ -421,6 +428,7 @@
                     depth=depth,
                     done=done,
                     print_type=print_type,
+                    print_shape=print_shape,
                     file=_file,
                     id_type=id_type,
                     stop_on_name=stop_on_name,
@@ -452,6 +460,7 @@
     depth: int = -1,
     done: dict[Literal["output"] | Variable | Apply, str] | None = None,
     print_type: bool = False,
+    print_shape: bool = False,
     file: TextIO = sys.stdout,
     print_destroy_map: bool = False,
     print_view_map: bool = False,
@@ -484,6 +493,8 @@
         See `debugprint`.
     print_type
         See `debugprint`.
+    print_shape
+        See `debugprint`.
     file
         File-like object to which to print.
     print_destroy_map
@@ -532,6 +543,11 @@
     else:
         type_str = ""
 
+    if print_shape:
+        shape_str = f" shape={str(var.type.shape).replace('None', '?')}"
+    else:
+        shape_str = ""
+
     if prefix_child is None:
         prefix_child = prefix
 
@@ -612,7 +628,7 @@
         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}"
+            var_output = f"{prefix}{node.op}{output_idx}{id_str}{type_str}{shape_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))
@@ -662,6 +678,7 @@
                     depth=depth - 1,
                     done=_done,
                     print_type=print_type,
+                    print_shape=print_shape,
                     file=file,
                     topo_order=topo_order,
                     id_type=id_type,
@@ -692,7 +709,7 @@
         else:
             data = ""
 
-        var_output = f"{prefix}{var}{id_str}{type_str}{data}"
+        var_output = f"{prefix}{var}{id_str}{type_str}{shape_str}{data}"
 
         if print_op_info and var.owner and var.owner not in op_information:
             op_information.update(op_debug_information(var.owner.op, var.owner))

pytensor/tensor/type.py

@@ -399,22 +399,13 @@ def __str__(self):
         else:
             shape = self.shape
             len_shape = len(shape)
-
-            def shape_str(s):
-                if s is None:
-                    return "?"
-                else:
-                    return str(s)
-
-            formatted_shape = ", ".join(shape_str(s) for s in shape)
-            if len_shape == 1:
-                formatted_shape += ","
+            formatted_shape = str(shape).replace("None", "?")
 
             if len_shape > 2:
                 name = f"Tensor{len_shape}"
             else:
                 name = ("Scalar", "Vector", "Matrix")[len_shape]
-            return f"{name}({self.dtype}, shape=({formatted_shape}))"
+            return f"{name}({self.dtype}, shape={formatted_shape})"
 
     def __repr__(self):
         return f"TensorType({self.dtype}, shape={self.shape})"