Skip to content

Commit 640c759

Browse files
committed
Add print_memory_map option to debugprint to enable destroy and view maps
1 parent 9451251 commit 640c759

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

pytensor/printing.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ def debugprint(
9999
print_op_info: bool = False,
100100
print_destroy_map: bool = False,
101101
print_view_map: bool = False,
102+
print_memory_map: bool = False,
102103
print_fgraph_inputs: bool = False,
103104
) -> str | TextIO:
104105
r"""Print a graph as text.
@@ -156,6 +157,8 @@ def debugprint(
156157
Whether to print the `destroy_map`\s of printed objects
157158
print_view_map
158159
Whether to print the `view_map`\s of printed objects
160+
print_memory_map
161+
Whether to set both `print_destroy_map` and `print_view_map` to ``True``.
159162
print_fgraph_inputs
160163
Print the inputs of `FunctionGraph`\s.
161164
@@ -180,6 +183,10 @@ def debugprint(
180183
if used_ids is None:
181184
used_ids = dict()
182185

186+
if print_memory_map:
187+
print_destroy_map = True
188+
print_view_map = True
189+
183190
inputs_to_print = []
184191
outputs_to_print = []
185192
profile_list: list[Any | None] = []

0 commit comments

Comments
 (0)