Skip to content

P
pytensor
提交 1f2542eb authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者： Ricardo Vieira
浏览文件
操作

Improve debug_print formatting

* Use better ASCII symbols * Show continuation hint for already seen nodes * Don't repeat entries for multiple output inner graphs
上级 39c1b49b
隐藏空白字符变更
内嵌 并排
正在显示 包含 540 行增加 和 485 行删除
pytensor/printing.py
浏览文件 @ 1f2542eb
...@@ -291,7 +291,7 @@ N.B.: ...@@ -291,7 +291,7 @@ N.B.:
for var in inputs_to_print: for var in inputs_to_print:
_debugprint( _debugprint(
var, var,
prefix="-", prefix="→ ",
depth=depth, depth=depth,
done=done, done=done,
print_type=print_type, print_type=print_type,
...@@ -342,11 +342,17 @@ N.B.: ...@@ -342,11 +342,17 @@ N.B.:
if len(inner_graph_vars) > 0: if len(inner_graph_vars) > 0:
print("", file=_file) print("", file=_file)
new_prefix = " >" prefix = ""
new_prefix_child = " >" new_prefix = prefix + " ← "
new_prefix_child = prefix + " "
print("Inner graphs:", file=_file) print("Inner graphs:", file=_file)
printed_inner_graphs_nodes = set()
for ig_var in inner_graph_vars: 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 # This is a work-around to maintain backward compatibility
# (e.g. to only print inner graphs that have been compiled through # (e.g. to only print inner graphs that have been compiled through
# a call to `Op.prepare_node`) # a call to `Op.prepare_node`)
...@@ -385,6 +391,7 @@ N.B.: ...@@ -385,6 +391,7 @@ N.B.:
_debugprint( _debugprint(
ig_var, ig_var,
prefix=prefix,
depth=depth, depth=depth,
done=done, done=done,
print_type=print_type, print_type=print_type,
...@@ -399,13 +406,14 @@ N.B.: ...@@ -399,13 +406,14 @@ N.B.:
print_op_info=print_op_info, print_op_info=print_op_info,
print_destroy_map=print_destroy_map, print_destroy_map=print_destroy_map,
print_view_map=print_view_map, print_view_map=print_view_map,
is_inner_graph_header=True,
) )
if print_fgraph_inputs: if print_fgraph_inputs:
for inp in inner_inputs: for inp in inner_inputs:
_debugprint( _debugprint(
inp, inp,
prefix="-", prefix=" → ",
depth=depth, depth=depth,
done=done, done=done,
print_type=print_type, print_type=print_type,
...@@ -485,6 +493,7 @@ def _debugprint( ...@@ -485,6 +493,7 @@ def _debugprint(
parent_node: Optional[Apply] = None, parent_node: Optional[Apply] = None,
print_op_info: bool = False, print_op_info: bool = False,
inner_graph_node: Optional[Apply] = None, inner_graph_node: Optional[Apply] = None,
is_inner_graph_header: bool = False,
) -> TextIO: ) -> TextIO:
r"""Print the graph represented by `var`. r"""Print the graph represented by `var`.
...@@ -625,7 +634,10 @@ def _debugprint( ...@@ -625,7 +634,10 @@ def _debugprint(
else: else:
data = "" 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: if print_op_info and node not in op_information:
op_information.update(op_debug_information(node.op, node)) op_information.update(op_debug_information(node.op, node))
...@@ -633,7 +645,7 @@ def _debugprint( ...@@ -633,7 +645,7 @@ def _debugprint(
node_info = ( node_info = (
parent_node and op_information.get(parent_node) parent_node and op_information.get(parent_node)
) or op_information.get(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]})" var_output = f"{var_output} ({node_info[var]})"
if profile and profile.apply_time and node in profile.apply_time: if profile and profile.apply_time and node in profile.apply_time:
...@@ -660,12 +672,13 @@ def _debugprint( ...@@ -660,12 +672,13 @@ def _debugprint(
if not already_done and ( if not already_done and (
not stop_on_name or not (hasattr(var, "name") and var.name is not None) not stop_on_name or not (hasattr(var, "name") and var.name is not None)
): ):
new_prefix = prefix_child + " |" new_prefix = prefix_child + " ├─ "
new_prefix_child = prefix_child + " |" new_prefix_child = prefix_child + " │ "
for in_idx, in_var in enumerate(node.inputs): for in_idx, in_var in enumerate(node.inputs):
if in_idx == len(node.inputs) - 1: 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 hasattr(in_var, "owner") and hasattr(in_var.owner, "op"):
if ( if (
...@@ -698,6 +711,8 @@ def _debugprint( ...@@ -698,6 +711,8 @@ def _debugprint(
print_view_map=print_view_map, print_view_map=print_view_map,
inner_graph_node=inner_graph_node, inner_graph_node=inner_graph_node,
) )
elif not is_inner_graph_header:
print(prefix_child + " └─ ···", file=file)
else: else:
id_str = get_id_str(var) id_str = get_id_str(var)
......
tests/compile/test_builders.py
浏览文件 @ 1f2542eb
...@@ -572,18 +572,18 @@ def test_debugprint(): ...@@ -572,18 +572,18 @@ def test_debugprint():
lines = output_str.split("\n") lines = output_str.split("\n")
exp_res = """OpFromGraph{inline=False} [id A] exp_res = """OpFromGraph{inline=False} [id A]
|x [id B] ├─ x [id B]
|y [id C] ├─ y [id C]
|z [id D] └─ z [id D]
Inner graphs: Inner graphs:
OpFromGraph{inline=False} [id A] OpFromGraph{inline=False} [id A]
>Elemwise{add,no_inplace} [id E] ← Elemwise{add,no_inplace} [id E]
> |*0-<TensorType(float64, (?, ?))> [id F] ├─ *0-<TensorType(float64, (?, ?))> [id F]
> |Elemwise{mul,no_inplace} [id G] └─ Elemwise{mul,no_inplace} [id G]
> |*1-<TensorType(float64, (?, ?))> [id H] ├─ *1-<TensorType(float64, (?, ?))> [id H]
> |*2-<TensorType(float64, (?, ?))> [id I] └─ *2-<TensorType(float64, (?, ?))> [id I]
""" """
for truth, out in zip(exp_res.split("\n"), lines): for truth, out in zip(exp_res.split("\n"), lines):
......
tests/scan/test_printing.py
浏览文件 @ 1f2542eb
...@@ -27,39 +27,42 @@ def test_debugprint_sitsot(): ...@@ -27,39 +27,42 @@ def test_debugprint_sitsot():
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] expected_output = """Subtensor{int64} [id A]
|Subtensor{int64::} [id B] ├─ Subtensor{int64::} [id B]
| |for{cpu,scan_fn} [id C] (outer_out_sit_sot-0) │ ├─ for{cpu,scan_fn} [id C] (outer_out_sit_sot-0)
| | |k [id D] (n_steps) │ │ ├─ k [id D] (n_steps)
| | |IncSubtensor{Set;:int64:} [id E] (outer_in_sit_sot-0) │ │ ├─ IncSubtensor{Set;:int64:} [id E] (outer_in_sit_sot-0)
| | | |AllocEmpty{dtype='float64'} [id F] │ │ │ ├─ AllocEmpty{dtype='float64'} [id F]
| | | | |Elemwise{add,no_inplace} [id G] │ │ │ │ ├─ Elemwise{add,no_inplace} [id G]
| | | | | |k [id D] │ │ │ │ │ ├─ k [id D]
| | | | | |Subtensor{int64} [id H] │ │ │ │ │ └─ Subtensor{int64} [id H]
| | | | | |Shape [id I] │ │ │ │ │ ├─ Shape [id I]
| | | | | | |Unbroadcast{0} [id J] │ │ │ │ │ │ └─ Unbroadcast{0} [id J]
| | | | | | |InplaceDimShuffle{x,0} [id K] │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id K]
| | | | | | |Elemwise{second,no_inplace} [id L] │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id L]
| | | | | | |A [id M] │ │ │ │ │ │ ├─ A [id M]
| | | | | | |InplaceDimShuffle{x} [id N] │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id N]
| | | | | | |TensorConstant{1.0} [id O] │ │ │ │ │ │ └─ TensorConstant{1.0} [id O]
| | | | | |ScalarConstant{0} [id P] │ │ │ │ │ └─ ScalarConstant{0} [id P]
| | | | |Subtensor{int64} [id Q] │ │ │ │ └─ Subtensor{int64} [id Q]
| | | | |Shape [id R] │ │ │ │ ├─ Shape [id R]
| | | | | |Unbroadcast{0} [id J] │ │ │ │ │ └─ Unbroadcast{0} [id J]
| | | | |ScalarConstant{1} [id S] │ │ │ │ │ └─ ···
| | | |Unbroadcast{0} [id J] │ │ │ │ └─ ScalarConstant{1} [id S]
| | | |ScalarFromTensor [id T] │ │ │ ├─ Unbroadcast{0} [id J]
| | | |Subtensor{int64} [id H] │ │ │ │ └─ ···
| | |A [id M] (outer_in_non_seqs-0) │ │ │ └─ ScalarFromTensor [id T]
| |ScalarConstant{1} [id U] │ │ │ └─ Subtensor{int64} [id H]
|ScalarConstant{-1} [id V] │ │ │ └─ ···
│ │ └─ A [id M] (outer_in_non_seqs-0)
│ └─ ScalarConstant{1} [id U]
└─ ScalarConstant{-1} [id V]
Inner graphs: Inner graphs:
for{cpu,scan_fn} [id C] (outer_out_sit_sot-0) for{cpu,scan_fn} [id C]
>Elemwise{mul,no_inplace} [id W] (inner_out_sit_sot-0) ← Elemwise{mul,no_inplace} [id W] (inner_out_sit_sot-0)
> |*0-<TensorType(float64, (?,))> [id X] -> [id E] (inner_in_sit_sot-0) ├─ *0-<TensorType(float64, (?,))> [id X] -> [id E] (inner_in_sit_sot-0)
> |*1-<TensorType(float64, (?,))> [id Y] -> [id M] (inner_in_non_seqs-0)""" └─ *1-<TensorType(float64, (?,))> [id Y] -> [id M] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -82,39 +85,42 @@ def test_debugprint_sitsot_no_extra_info(): ...@@ -82,39 +85,42 @@ def test_debugprint_sitsot_no_extra_info():
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] expected_output = """Subtensor{int64} [id A]
|Subtensor{int64::} [id B] ├─ Subtensor{int64::} [id B]
| |for{cpu,scan_fn} [id C] │ ├─ for{cpu,scan_fn} [id C]
| | |k [id D] │ │ ├─ k [id D]
| | |IncSubtensor{Set;:int64:} [id E] │ │ ├─ IncSubtensor{Set;:int64:} [id E]
| | | |AllocEmpty{dtype='float64'} [id F] │ │ │ ├─ AllocEmpty{dtype='float64'} [id F]
| | | | |Elemwise{add,no_inplace} [id G] │ │ │ │ ├─ Elemwise{add,no_inplace} [id G]
| | | | | |k [id D] │ │ │ │ │ ├─ k [id D]
| | | | | |Subtensor{int64} [id H] │ │ │ │ │ └─ Subtensor{int64} [id H]
| | | | | |Shape [id I] │ │ │ │ │ ├─ Shape [id I]
| | | | | | |Unbroadcast{0} [id J] │ │ │ │ │ │ └─ Unbroadcast{0} [id J]
| | | | | | |InplaceDimShuffle{x,0} [id K] │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id K]
| | | | | | |Elemwise{second,no_inplace} [id L] │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id L]
| | | | | | |A [id M] │ │ │ │ │ │ ├─ A [id M]
| | | | | | |InplaceDimShuffle{x} [id N] │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id N]
| | | | | | |TensorConstant{1.0} [id O] │ │ │ │ │ │ └─ TensorConstant{1.0} [id O]
| | | | | |ScalarConstant{0} [id P] │ │ │ │ │ └─ ScalarConstant{0} [id P]
| | | | |Subtensor{int64} [id Q] │ │ │ │ └─ Subtensor{int64} [id Q]
| | | | |Shape [id R] │ │ │ │ ├─ Shape [id R]
| | | | | |Unbroadcast{0} [id J] │ │ │ │ │ └─ Unbroadcast{0} [id J]
| | | | |ScalarConstant{1} [id S] │ │ │ │ │ └─ ···
| | | |Unbroadcast{0} [id J] │ │ │ │ └─ ScalarConstant{1} [id S]
| | | |ScalarFromTensor [id T] │ │ │ ├─ Unbroadcast{0} [id J]
| | | |Subtensor{int64} [id H] │ │ │ │ └─ ···
| | |A [id M] │ │ │ └─ ScalarFromTensor [id T]
| |ScalarConstant{1} [id U] │ │ │ └─ Subtensor{int64} [id H]
|ScalarConstant{-1} [id V] │ │ │ └─ ···
│ │ └─ A [id M]
│ └─ ScalarConstant{1} [id U]
└─ ScalarConstant{-1} [id V]
Inner graphs: Inner graphs:
for{cpu,scan_fn} [id C] for{cpu,scan_fn} [id C]
>Elemwise{mul,no_inplace} [id W] ← Elemwise{mul,no_inplace} [id W]
> |*0-<TensorType(float64, (?,))> [id X] -> [id E] ├─ *0-<TensorType(float64, (?,))> [id X] -> [id E]
> |*1-<TensorType(float64, (?,))> [id Y] -> [id M]""" └─ *1-<TensorType(float64, (?,))> [id Y] -> [id M]"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -142,42 +148,47 @@ def test_debugprint_nitsot(): ...@@ -142,42 +148,47 @@ def test_debugprint_nitsot():
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A] expected_output = """Sum{acc_dtype=float64} [id A]
|for{cpu,scan_fn} [id B] (outer_out_nit_sot-0) └─ for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
|Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0) ├─ Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
| |Subtensor{int64} [id D] │ ├─ Subtensor{int64} [id D]
| | |Shape [id E] │ │ ├─ Shape [id E]
| | | |Subtensor{int64::} [id F] 'coefficients[0:]' │ │ │ └─ Subtensor{int64::} [id F] 'coefficients[0:]'
| | | |coefficients [id G] │ │ │ ├─ coefficients [id G]
| | | |ScalarConstant{0} [id H] │ │ │ └─ ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I] │ │ └─ ScalarConstant{0} [id I]
| |Subtensor{int64} [id J] │ └─ Subtensor{int64} [id J]
| |Shape [id K] │ ├─ Shape [id K]
| | |Subtensor{int64::} [id L] │ │ └─ Subtensor{int64::} [id L]
| | |ARange{dtype='int64'} [id M] │ │ ├─ ARange{dtype='int64'} [id M]
| | | |TensorConstant{0} [id N] │ │ │ ├─ TensorConstant{0} [id N]
| | | |TensorConstant{10000} [id O] │ │ │ ├─ TensorConstant{10000} [id O]
| | | |TensorConstant{1} [id P] │ │ │ └─ TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q] │ │ └─ ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R] │ └─ ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] (outer_in_seqs-0) ├─ Subtensor{:int64:} [id S] (outer_in_seqs-0)
| |Subtensor{int64::} [id F] 'coefficients[0:]' │ ├─ Subtensor{int64::} [id F] 'coefficients[0:]'
| |ScalarFromTensor [id T] │ │ └─ ···
| |Elemwise{scalar_minimum,no_inplace} [id C] │ └─ ScalarFromTensor [id T]
|Subtensor{:int64:} [id U] (outer_in_seqs-1) │ └─ Elemwise{scalar_minimum,no_inplace} [id C]
| |Subtensor{int64::} [id L] │ └─ ···
| |ScalarFromTensor [id V] ├─ Subtensor{:int64:} [id U] (outer_in_seqs-1)
| |Elemwise{scalar_minimum,no_inplace} [id C] │ ├─ Subtensor{int64::} [id L]
|Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0) │ │ └─ ···
|x [id W] (outer_in_non_seqs-0) │ └─ ScalarFromTensor [id V]
│ └─ Elemwise{scalar_minimum,no_inplace} [id C]
│ └─ ···
├─ Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
│ └─ ···
└─ x [id W] (outer_in_non_seqs-0)
Inner graphs: Inner graphs:
for{cpu,scan_fn} [id B] (outer_out_nit_sot-0) for{cpu,scan_fn} [id B]
>Elemwise{mul,no_inplace} [id X] (inner_out_nit_sot-0) ← Elemwise{mul,no_inplace} [id X] (inner_out_nit_sot-0)
> |*0-<TensorType(float64, ())> [id Y] -> [id S] (inner_in_seqs-0) ├─ *0-<TensorType(float64, ())> [id Y] -> [id S] (inner_in_seqs-0)
> |Elemwise{pow,no_inplace} [id Z] └─ Elemwise{pow,no_inplace} [id Z]
> |*2-<TensorType(float64, ())> [id BA] -> [id W] (inner_in_non_seqs-0) ├─ *2-<TensorType(float64, ())> [id BA] -> [id W] (inner_in_non_seqs-0)
> |*1-<TensorType(int64, ())> [id BB] -> [id U] (inner_in_seqs-1)""" └─ *1-<TensorType(int64, ())> [id BB] -> [id U] (inner_in_seqs-1)"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -215,76 +226,84 @@ def test_debugprint_nested_scans(): ...@@ -215,76 +226,84 @@ def test_debugprint_nested_scans():
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A] expected_output = """Sum{acc_dtype=float64} [id A]
|for{cpu,scan_fn} [id B] (outer_out_nit_sot-0) └─ for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
|Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0) ├─ Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
| |Subtensor{int64} [id D] │ ├─ Subtensor{int64} [id D]
| | |Shape [id E] │ │ ├─ Shape [id E]
| | | |Subtensor{int64::} [id F] 'c[0:]' │ │ │ └─ Subtensor{int64::} [id F] 'c[0:]'
| | | |c [id G] │ │ │ ├─ c [id G]
| | | |ScalarConstant{0} [id H] │ │ │ └─ ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I] │ │ └─ ScalarConstant{0} [id I]
| |Subtensor{int64} [id J] │ └─ Subtensor{int64} [id J]
| |Shape [id K] │ ├─ Shape [id K]
| | |Subtensor{int64::} [id L] │ │ └─ Subtensor{int64::} [id L]
| | |ARange{dtype='int64'} [id M] │ │ ├─ ARange{dtype='int64'} [id M]
| | | |TensorConstant{0} [id N] │ │ │ ├─ TensorConstant{0} [id N]
| | | |TensorConstant{10} [id O] │ │ │ ├─ TensorConstant{10} [id O]
| | | |TensorConstant{1} [id P] │ │ │ └─ TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q] │ │ └─ ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R] │ └─ ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] (outer_in_seqs-0) ├─ Subtensor{:int64:} [id S] (outer_in_seqs-0)
| |Subtensor{int64::} [id F] 'c[0:]' │ ├─ Subtensor{int64::} [id F] 'c[0:]'
| |ScalarFromTensor [id T] │ │ └─ ···
| |Elemwise{scalar_minimum,no_inplace} [id C] │ └─ ScalarFromTensor [id T]
|Subtensor{:int64:} [id U] (outer_in_seqs-1) │ └─ Elemwise{scalar_minimum,no_inplace} [id C]
| |Subtensor{int64::} [id L] │ └─ ···
| |ScalarFromTensor [id V] ├─ Subtensor{:int64:} [id U] (outer_in_seqs-1)
| |Elemwise{scalar_minimum,no_inplace} [id C] │ ├─ Subtensor{int64::} [id L]
|Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0) │ │ └─ ···
|A [id W] (outer_in_non_seqs-0) │ └─ ScalarFromTensor [id V]
|k [id X] (outer_in_non_seqs-1) │ └─ Elemwise{scalar_minimum,no_inplace} [id C]
│ └─ ···
├─ Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
│ └─ ···
├─ A [id W] (outer_in_non_seqs-0)
└─ k [id X] (outer_in_non_seqs-1)
Inner graphs: Inner graphs:
for{cpu,scan_fn} [id B] (outer_out_nit_sot-0) for{cpu,scan_fn} [id B]
>Elemwise{mul,no_inplace} [id Y] (inner_out_nit_sot-0) ← Elemwise{mul,no_inplace} [id Y] (inner_out_nit_sot-0)
> |InplaceDimShuffle{x} [id Z] ├─ InplaceDimShuffle{x} [id Z]
> | |*0-<TensorType(float64, ())> [id BA] -> [id S] (inner_in_seqs-0) │ └─ *0-<TensorType(float64, ())> [id BA] -> [id S] (inner_in_seqs-0)
> |Elemwise{pow,no_inplace} [id BB] └─ Elemwise{pow,no_inplace} [id BB]
> |Subtensor{int64} [id BC] ├─ Subtensor{int64} [id BC]
> | |Subtensor{int64::} [id BD] │ ├─ Subtensor{int64::} [id BD]
> | | |for{cpu,scan_fn} [id BE] (outer_out_sit_sot-0) │ │ ├─ for{cpu,scan_fn} [id BE] (outer_out_sit_sot-0)
> | | | |*3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1) (n_steps) │ │ │ ├─ *3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1) (n_steps)
> | | | |IncSubtensor{Set;:int64:} [id BG] (outer_in_sit_sot-0) │ │ │ ├─ IncSubtensor{Set;:int64:} [id BG] (outer_in_sit_sot-0)
> | | | | |AllocEmpty{dtype='float64'} [id BH] │ │ │ │ ├─ AllocEmpty{dtype='float64'} [id BH]
> | | | | | |Elemwise{add,no_inplace} [id BI] │ │ │ │ │ ├─ Elemwise{add,no_inplace} [id BI]
> | | | | | | |*3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1) │ │ │ │ │ │ ├─ *3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1)
> | | | | | | |Subtensor{int64} [id BJ] │ │ │ │ │ │ └─ Subtensor{int64} [id BJ]
> | | | | | | |Shape [id BK] │ │ │ │ │ │ ├─ Shape [id BK]
> | | | | | | | |Unbroadcast{0} [id BL] │ │ │ │ │ │ │ └─ Unbroadcast{0} [id BL]
> | | | | | | | |InplaceDimShuffle{x,0} [id BM] │ │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id BM]
> | | | | | | | |Elemwise{second,no_inplace} [id BN] │ │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id BN]
> | | | | | | | |*2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0) │ │ │ │ │ │ │ ├─ *2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0)
> | | | | | | | |InplaceDimShuffle{x} [id BP] │ │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id BP]
> | | | | | | | |TensorConstant{1.0} [id BQ] │ │ │ │ │ │ │ └─ TensorConstant{1.0} [id BQ]
> | | | | | | |ScalarConstant{0} [id BR] │ │ │ │ │ │ └─ ScalarConstant{0} [id BR]
> | | | | | |Subtensor{int64} [id BS] │ │ │ │ │ └─ Subtensor{int64} [id BS]
> | | | | | |Shape [id BT] │ │ │ │ │ ├─ Shape [id BT]
> | | | | | | |Unbroadcast{0} [id BL] │ │ │ │ │ │ └─ Unbroadcast{0} [id BL]
> | | | | | |ScalarConstant{1} [id BU] │ │ │ │ │ │ └─ ···
> | | | | |Unbroadcast{0} [id BL] │ │ │ │ │ └─ ScalarConstant{1} [id BU]
> | | | | |ScalarFromTensor [id BV] │ │ │ │ ├─ Unbroadcast{0} [id BL]
> | | | | |Subtensor{int64} [id BJ] │ │ │ │ │ └─ ···
> | | | |*2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0) (outer_in_non_seqs-0) │ │ │ │ └─ ScalarFromTensor [id BV]
> | | |ScalarConstant{1} [id BW] │ │ │ │ └─ Subtensor{int64} [id BJ]
> | |ScalarConstant{-1} [id BX] │ │ │ │ └─ ···
> |InplaceDimShuffle{x} [id BY] │ │ │ └─ *2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
> |*1-<TensorType(int64, ())> [id BZ] -> [id U] (inner_in_seqs-1) │ │ └─ ScalarConstant{1} [id BW]
│ └─ ScalarConstant{-1} [id BX]
for{cpu,scan_fn} [id BE] (outer_out_sit_sot-0) └─ InplaceDimShuffle{x} [id BY]
>Elemwise{mul,no_inplace} [id CA] (inner_out_sit_sot-0) └─ *1-<TensorType(int64, ())> [id BZ] -> [id U] (inner_in_seqs-1)
> |*0-<TensorType(float64, (?,))> [id CB] -> [id BG] (inner_in_sit_sot-0)
> |*1-<TensorType(float64, (?,))> [id CC] -> [id BO] (inner_in_non_seqs-0)""" for{cpu,scan_fn} [id BE]
← Elemwise{mul,no_inplace} [id CA] (inner_out_sit_sot-0)
├─ *0-<TensorType(float64, (?,))> [id CB] -> [id BG] (inner_in_sit_sot-0)
└─ *1-<TensorType(float64, (?,))> [id CC] -> [id BO] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -296,86 +315,94 @@ def test_debugprint_nested_scans(): ...@@ -296,86 +315,94 @@ def test_debugprint_nested_scans():
) )
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """-c [id A] expected_output = """→ c [id A]
-k [id B] → k [id B]
-A [id C] → A [id C]
Sum{acc_dtype=float64} [id D] 13 Sum{acc_dtype=float64} [id D] 13
|for{cpu,scan_fn} [id E] 12 (outer_out_nit_sot-0) └─ for{cpu,scan_fn} [id E] 12 (outer_out_nit_sot-0)
|Elemwise{scalar_minimum,no_inplace} [id F] 7 (outer_in_nit_sot-0) ├─ Elemwise{scalar_minimum,no_inplace} [id F] 7 (outer_in_nit_sot-0)
| |Subtensor{int64} [id G] 6 │ ├─ Subtensor{int64} [id G] 6
| | |Shape [id H] 5 │ │ ├─ Shape [id H] 5
| | | |Subtensor{int64::} [id I] 'c[0:]' 4 │ │ │ └─ Subtensor{int64::} [id I] 'c[0:]' 4
| | | |c [id A] │ │ │ ├─ c [id A]
| | | |ScalarConstant{0} [id J] │ │ │ └─ ScalarConstant{0} [id J]
| | |ScalarConstant{0} [id K] │ │ └─ ScalarConstant{0} [id K]
| |Subtensor{int64} [id L] 3 │ └─ Subtensor{int64} [id L] 3
| |Shape [id M] 2 │ ├─ Shape [id M] 2
| | |Subtensor{int64::} [id N] 1 │ │ └─ Subtensor{int64::} [id N] 1
| | |ARange{dtype='int64'} [id O] 0 │ │ ├─ ARange{dtype='int64'} [id O] 0
| | | |TensorConstant{0} [id P] │ │ │ ├─ TensorConstant{0} [id P]
| | | |TensorConstant{10} [id Q] │ │ │ ├─ TensorConstant{10} [id Q]
| | | |TensorConstant{1} [id R] │ │ │ └─ TensorConstant{1} [id R]
| | |ScalarConstant{0} [id S] │ │ └─ ScalarConstant{0} [id S]
| |ScalarConstant{0} [id T] │ └─ ScalarConstant{0} [id T]
|Subtensor{:int64:} [id U] 11 (outer_in_seqs-0) ├─ Subtensor{:int64:} [id U] 11 (outer_in_seqs-0)
| |Subtensor{int64::} [id I] 'c[0:]' 4 │ ├─ Subtensor{int64::} [id I] 'c[0:]' 4
| |ScalarFromTensor [id V] 10 │ │ └─ ···
| |Elemwise{scalar_minimum,no_inplace} [id F] 7 │ └─ ScalarFromTensor [id V] 10
|Subtensor{:int64:} [id W] 9 (outer_in_seqs-1) │ └─ Elemwise{scalar_minimum,no_inplace} [id F] 7
| |Subtensor{int64::} [id N] 1 │ └─ ···
| |ScalarFromTensor [id X] 8 ├─ Subtensor{:int64:} [id W] 9 (outer_in_seqs-1)
| |Elemwise{scalar_minimum,no_inplace} [id F] 7 │ ├─ Subtensor{int64::} [id N] 1
|Elemwise{scalar_minimum,no_inplace} [id F] 7 (outer_in_nit_sot-0) │ │ └─ ···
|A [id C] (outer_in_non_seqs-0) │ └─ ScalarFromTensor [id X] 8
|k [id B] (outer_in_non_seqs-1) │ └─ Elemwise{scalar_minimum,no_inplace} [id F] 7
│ └─ ···
├─ Elemwise{scalar_minimum,no_inplace} [id F] 7 (outer_in_nit_sot-0)
│ └─ ···
├─ A [id C] (outer_in_non_seqs-0)
└─ k [id B] (outer_in_non_seqs-1)
Inner graphs: Inner graphs:
for{cpu,scan_fn} [id E] (outer_out_nit_sot-0) for{cpu,scan_fn} [id E]
-*0-<TensorType(float64, ())> [id Y] -> [id U] (inner_in_seqs-0) → *0-<TensorType(float64, ())> [id Y] -> [id U] (inner_in_seqs-0)
-*1-<TensorType(int64, ())> [id Z] -> [id W] (inner_in_seqs-1) → *1-<TensorType(int64, ())> [id Z] -> [id W] (inner_in_seqs-1)
-*2-<TensorType(float64, (?,))> [id BA] -> [id C] (inner_in_non_seqs-0) → *2-<TensorType(float64, (?,))> [id BA] -> [id C] (inner_in_non_seqs-0)
-*3-<TensorType(int32, ())> [id BB] -> [id B] (inner_in_non_seqs-1) → *3-<TensorType(int32, ())> [id BB] -> [id B] (inner_in_non_seqs-1)
>Elemwise{mul,no_inplace} [id BC] (inner_out_nit_sot-0) ← Elemwise{mul,no_inplace} [id BC] (inner_out_nit_sot-0)
> |InplaceDimShuffle{x} [id BD] ├─ InplaceDimShuffle{x} [id BD]
> | |*0-<TensorType(float64, ())> [id Y] (inner_in_seqs-0) │ └─ *0-<TensorType(float64, ())> [id Y] (inner_in_seqs-0)
> |Elemwise{pow,no_inplace} [id BE] └─ Elemwise{pow,no_inplace} [id BE]
> |Subtensor{int64} [id BF] ├─ Subtensor{int64} [id BF]
> | |Subtensor{int64::} [id BG] │ ├─ Subtensor{int64::} [id BG]
> | | |for{cpu,scan_fn} [id BH] (outer_out_sit_sot-0) │ │ ├─ for{cpu,scan_fn} [id BH] (outer_out_sit_sot-0)
> | | | |*3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1) (n_steps) │ │ │ ├─ *3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1) (n_steps)
> | | | |IncSubtensor{Set;:int64:} [id BI] (outer_in_sit_sot-0) │ │ │ ├─ IncSubtensor{Set;:int64:} [id BI] (outer_in_sit_sot-0)
> | | | | |AllocEmpty{dtype='float64'} [id BJ] │ │ │ │ ├─ AllocEmpty{dtype='float64'} [id BJ]
> | | | | | |Elemwise{add,no_inplace} [id BK] │ │ │ │ │ ├─ Elemwise{add,no_inplace} [id BK]
> | | | | | | |*3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1) │ │ │ │ │ │ ├─ *3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1)
> | | | | | | |Subtensor{int64} [id BL] │ │ │ │ │ │ └─ Subtensor{int64} [id BL]
> | | | | | | |Shape [id BM] │ │ │ │ │ │ ├─ Shape [id BM]
> | | | | | | | |Unbroadcast{0} [id BN] │ │ │ │ │ │ │ └─ Unbroadcast{0} [id BN]
> | | | | | | | |InplaceDimShuffle{x,0} [id BO] │ │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id BO]
> | | | | | | | |Elemwise{second,no_inplace} [id BP] │ │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id BP]
> | | | | | | | |*2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0) │ │ │ │ │ │ │ ├─ *2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0)
> | | | | | | | |InplaceDimShuffle{x} [id BQ] │ │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id BQ]
> | | | | | | | |TensorConstant{1.0} [id BR] │ │ │ │ │ │ │ └─ TensorConstant{1.0} [id BR]
> | | | | | | |ScalarConstant{0} [id BS] │ │ │ │ │ │ └─ ScalarConstant{0} [id BS]
> | | | | | |Subtensor{int64} [id BT] │ │ │ │ │ └─ Subtensor{int64} [id BT]
> | | | | | |Shape [id BU] │ │ │ │ │ ├─ Shape [id BU]
> | | | | | | |Unbroadcast{0} [id BN] │ │ │ │ │ │ └─ Unbroadcast{0} [id BN]
> | | | | | |ScalarConstant{1} [id BV] │ │ │ │ │ │ └─ ···
> | | | | |Unbroadcast{0} [id BN] │ │ │ │ │ └─ ScalarConstant{1} [id BV]
> | | | | |ScalarFromTensor [id BW] │ │ │ │ ├─ Unbroadcast{0} [id BN]
> | | | | |Subtensor{int64} [id BL] │ │ │ │ │ └─ ···
> | | | |*2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0) (outer_in_non_seqs-0) │ │ │ │ └─ ScalarFromTensor [id BW]
> | | |ScalarConstant{1} [id BX] │ │ │ │ └─ Subtensor{int64} [id BL]
> | |ScalarConstant{-1} [id BY] │ │ │ │ └─ ···
> |InplaceDimShuffle{x} [id BZ] │ │ │ └─ *2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
> |*1-<TensorType(int64, ())> [id Z] (inner_in_seqs-1) │ │ └─ ScalarConstant{1} [id BX]
│ └─ ScalarConstant{-1} [id BY]
for{cpu,scan_fn} [id BH] (outer_out_sit_sot-0) └─ InplaceDimShuffle{x} [id BZ]
-*0-<TensorType(float64, (?,))> [id CA] -> [id BI] (inner_in_sit_sot-0) └─ *1-<TensorType(int64, ())> [id Z] (inner_in_seqs-1)
-*1-<TensorType(float64, (?,))> [id CB] -> [id BA] (inner_in_non_seqs-0)
>Elemwise{mul,no_inplace} [id CC] (inner_out_sit_sot-0) for{cpu,scan_fn} [id BH]
> |*0-<TensorType(float64, (?,))> [id CA] (inner_in_sit_sot-0) → *0-<TensorType(float64, (?,))> [id CA] -> [id BI] (inner_in_sit_sot-0)
> |*1-<TensorType(float64, (?,))> [id CB] (inner_in_non_seqs-0)""" → *1-<TensorType(float64, (?,))> [id CB] -> [id BA] (inner_in_non_seqs-0)
← Elemwise{mul,no_inplace} [id CC] (inner_out_sit_sot-0)
├─ *0-<TensorType(float64, (?,))> [id CA] (inner_in_sit_sot-0)
└─ *1-<TensorType(float64, (?,))> [id CB] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -403,53 +430,54 @@ def test_debugprint_mitsot(): ...@@ -403,53 +430,54 @@ def test_debugprint_mitsot():
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """Elemwise{add,no_inplace} [id A] expected_output = """Elemwise{add,no_inplace} [id A]
|Subtensor{int64::} [id B] ├─ Subtensor{int64::} [id B]
| |for{cpu,scan_fn}.0 [id C] (outer_out_mit_sot-0) │ ├─ for{cpu,scan_fn}.0 [id C] (outer_out_mit_sot-0)
| | |TensorConstant{5} [id D] (n_steps) │ │ ├─ TensorConstant{5} [id D] (n_steps)
| | |IncSubtensor{Set;:int64:} [id E] (outer_in_mit_sot-0) │ │ ├─ IncSubtensor{Set;:int64:} [id E] (outer_in_mit_sot-0)
| | | |AllocEmpty{dtype='int64'} [id F] │ │ │ ├─ AllocEmpty{dtype='int64'} [id F]
| | | | |Elemwise{add,no_inplace} [id G] │ │ │ │ └─ Elemwise{add,no_inplace} [id G]
| | | | |TensorConstant{5} [id D] │ │ │ │ ├─ TensorConstant{5} [id D]
| | | | |Subtensor{int64} [id H] │ │ │ │ └─ Subtensor{int64} [id H]
| | | | |Shape [id I] │ │ │ │ ├─ Shape [id I]
| | | | | |Subtensor{:int64:} [id J] │ │ │ │ │ └─ Subtensor{:int64:} [id J]
| | | | | |<TensorType(int64, (?,))> [id K] │ │ │ │ │ ├─ <TensorType(int64, (?,))> [id K]
| | | | | |ScalarConstant{2} [id L] │ │ │ │ │ └─ ScalarConstant{2} [id L]
| | | | |ScalarConstant{0} [id M] │ │ │ │ └─ ScalarConstant{0} [id M]
| | | |Subtensor{:int64:} [id J] │ │ │ ├─ Subtensor{:int64:} [id J]
| | | |ScalarFromTensor [id N] │ │ │ │ └─ ···
| | | |Subtensor{int64} [id H] │ │ │ └─ ScalarFromTensor [id N]
| | |IncSubtensor{Set;:int64:} [id O] (outer_in_mit_sot-1) │ │ │ └─ Subtensor{int64} [id H]
| | |AllocEmpty{dtype='int64'} [id P] │ │ │ └─ ···
| | | |Elemwise{add,no_inplace} [id Q] │ │ └─ IncSubtensor{Set;:int64:} [id O] (outer_in_mit_sot-1)
| | | |TensorConstant{5} [id D] │ │ ├─ AllocEmpty{dtype='int64'} [id P]
| | | |Subtensor{int64} [id R] │ │ │ └─ Elemwise{add,no_inplace} [id Q]
| | | |Shape [id S] │ │ │ ├─ TensorConstant{5} [id D]
| | | | |Subtensor{:int64:} [id T] │ │ │ └─ Subtensor{int64} [id R]
| | | | |<TensorType(int64, (?,))> [id U] │ │ │ ├─ Shape [id S]
| | | | |ScalarConstant{2} [id V] │ │ │ │ └─ Subtensor{:int64:} [id T]
| | | |ScalarConstant{0} [id W] │ │ │ │ ├─ <TensorType(int64, (?,))> [id U]
| | |Subtensor{:int64:} [id T] │ │ │ │ └─ ScalarConstant{2} [id V]
| | |ScalarFromTensor [id X] │ │ │ └─ ScalarConstant{0} [id W]
| | |Subtensor{int64} [id R] │ │ ├─ Subtensor{:int64:} [id T]
| |ScalarConstant{2} [id Y] │ │ │ └─ ···
|Subtensor{int64::} [id Z] │ │ └─ ScalarFromTensor [id X]
|for{cpu,scan_fn}.1 [id C] (outer_out_mit_sot-1) │ │ └─ Subtensor{int64} [id R]
|ScalarConstant{2} [id BA] │ │ └─ ···
│ └─ ScalarConstant{2} [id Y]
└─ Subtensor{int64::} [id Z]
├─ for{cpu,scan_fn}.1 [id C] (outer_out_mit_sot-1)
│ └─ ···
└─ ScalarConstant{2} [id BA]
Inner graphs: Inner graphs:
for{cpu,scan_fn}.0 [id C] (outer_out_mit_sot-0) for{cpu,scan_fn} [id C]
>Elemwise{add,no_inplace} [id BB] (inner_out_mit_sot-0) ← Elemwise{add,no_inplace} [id BB] (inner_out_mit_sot-0)
> |*1-<TensorType(int64, ())> [id BC] -> [id E] (inner_in_mit_sot-0-1) ├─ *1-<TensorType(int64, ())> [id BC] -> [id E] (inner_in_mit_sot-0-1)
> |*0-<TensorType(int64, ())> [id BD] -> [id E] (inner_in_mit_sot-0-0) └─ *0-<TensorType(int64, ())> [id BD] -> [id E] (inner_in_mit_sot-0-0)
>Elemwise{add,no_inplace} [id BE] (inner_out_mit_sot-1) ← Elemwise{add,no_inplace} [id BE] (inner_out_mit_sot-1)
> |*3-<TensorType(int64, ())> [id BF] -> [id O] (inner_in_mit_sot-1-1) ├─ *3-<TensorType(int64, ())> [id BF] -> [id O] (inner_in_mit_sot-1-1)
> |*2-<TensorType(int64, ())> [id BG] -> [id O] (inner_in_mit_sot-1-0) └─ *2-<TensorType(int64, ())> [id BG] -> [id O] (inner_in_mit_sot-1-0)"""
for{cpu,scan_fn}.1 [id C] (outer_out_mit_sot-1)
>Elemwise{add,no_inplace} [id BB] (inner_out_mit_sot-0)
>Elemwise{add,no_inplace} [id BE] (inner_out_mit_sot-1)"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -474,106 +502,118 @@ def test_debugprint_mitmot(): ...@@ -474,106 +502,118 @@ def test_debugprint_mitmot():
lines = output_str.split("\n") lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] expected_output = """Subtensor{int64} [id A]
|for{cpu,grad_of_scan_fn}.1 [id B] (outer_out_sit_sot-0) ├─ for{cpu,grad_of_scan_fn}.1 [id B] (outer_out_sit_sot-0)
| |Elemwise{sub,no_inplace} [id C] (n_steps) │ ├─ Elemwise{sub,no_inplace} [id C] (n_steps)
| | |Subtensor{int64} [id D] │ │ ├─ Subtensor{int64} [id D]
| | | |Shape [id E] │ │ │ ├─ Shape [id E]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0) │ │ │ │ └─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | |k [id G] (n_steps) │ │ │ │ ├─ k [id G] (n_steps)
| | | | |IncSubtensor{Set;:int64:} [id H] (outer_in_sit_sot-0) │ │ │ │ ├─ IncSubtensor{Set;:int64:} [id H] (outer_in_sit_sot-0)
| | | | | |AllocEmpty{dtype='float64'} [id I] │ │ │ │ │ ├─ AllocEmpty{dtype='float64'} [id I]
| | | | | | |Elemwise{add,no_inplace} [id J] │ │ │ │ │ │ ├─ Elemwise{add,no_inplace} [id J]
| | | | | | | |k [id G] │ │ │ │ │ │ │ ├─ k [id G]
| | | | | | | |Subtensor{int64} [id K] │ │ │ │ │ │ │ └─ Subtensor{int64} [id K]
| | | | | | | |Shape [id L] │ │ │ │ │ │ │ ├─ Shape [id L]
| | | | | | | | |Unbroadcast{0} [id M] │ │ │ │ │ │ │ │ └─ Unbroadcast{0} [id M]
| | | | | | | | |InplaceDimShuffle{x,0} [id N] │ │ │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id N]
| | | | | | | | |Elemwise{second,no_inplace} [id O] │ │ │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id O]
| | | | | | | | |A [id P] │ │ │ │ │ │ │ │ ├─ A [id P]
| | | | | | | | |InplaceDimShuffle{x} [id Q] │ │ │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id Q]
| | | | | | | | |TensorConstant{1.0} [id R] │ │ │ │ │ │ │ │ └─ TensorConstant{1.0} [id R]
| | | | | | | |ScalarConstant{0} [id S] │ │ │ │ │ │ │ └─ ScalarConstant{0} [id S]
| | | | | | |Subtensor{int64} [id T] │ │ │ │ │ │ └─ Subtensor{int64} [id T]
| | | | | | |Shape [id U] │ │ │ │ │ │ ├─ Shape [id U]
| | | | | | | |Unbroadcast{0} [id M] │ │ │ │ │ │ │ └─ Unbroadcast{0} [id M]
| | | | | | |ScalarConstant{1} [id V] │ │ │ │ │ │ │ └─ ···
| | | | | |Unbroadcast{0} [id M] │ │ │ │ │ │ └─ ScalarConstant{1} [id V]
| | | | | |ScalarFromTensor [id W] │ │ │ │ │ ├─ Unbroadcast{0} [id M]
| | | | | |Subtensor{int64} [id K] │ │ │ │ │ │ └─ ···
| | | | |A [id P] (outer_in_non_seqs-0) │ │ │ │ │ └─ ScalarFromTensor [id W]
| | | |ScalarConstant{0} [id X] │ │ │ │ │ └─ Subtensor{int64} [id K]
| | |TensorConstant{1} [id Y] │ │ │ │ │ └─ ···
| |Subtensor{:int64:} [id Z] (outer_in_seqs-0) │ │ │ │ └─ A [id P] (outer_in_non_seqs-0)
| | |Subtensor{::int64} [id BA] │ │ │ └─ ScalarConstant{0} [id X]
| | | |Subtensor{:int64:} [id BB] │ │ └─ TensorConstant{1} [id Y]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0) │ ├─ Subtensor{:int64:} [id Z] (outer_in_seqs-0)
| | | | |ScalarConstant{-1} [id BC] │ │ ├─ Subtensor{::int64} [id BA]
| | | |ScalarConstant{-1} [id BD] │ │ │ ├─ Subtensor{:int64:} [id BB]
| | |ScalarFromTensor [id BE] │ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | |Elemwise{sub,no_inplace} [id C] │ │ │ │ │ └─ ···
| |Subtensor{:int64:} [id BF] (outer_in_seqs-1) │ │ │ │ └─ ScalarConstant{-1} [id BC]
| | |Subtensor{:int64:} [id BG] │ │ │ └─ ScalarConstant{-1} [id BD]
| | | |Subtensor{::int64} [id BH] │ │ └─ ScalarFromTensor [id BE]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0) │ │ └─ Elemwise{sub,no_inplace} [id C]
| | | | |ScalarConstant{-1} [id BI] │ │ └─ ···
| | | |ScalarConstant{-1} [id BJ] │ ├─ Subtensor{:int64:} [id BF] (outer_in_seqs-1)
| | |ScalarFromTensor [id BK] │ │ ├─ Subtensor{:int64:} [id BG]
| | |Elemwise{sub,no_inplace} [id C] │ │ │ ├─ Subtensor{::int64} [id BH]
| |Subtensor{::int64} [id BL] (outer_in_mit_mot-0) │ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | |IncSubtensor{Inc;int64::} [id BM] │ │ │ │ │ └─ ···
| | | |Elemwise{second,no_inplace} [id BN] │ │ │ │ └─ ScalarConstant{-1} [id BI]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0) │ │ │ └─ ScalarConstant{-1} [id BJ]
| | | | |InplaceDimShuffle{x,x} [id BO] │ │ └─ ScalarFromTensor [id BK]
| | | | |TensorConstant{0.0} [id BP] │ │ └─ Elemwise{sub,no_inplace} [id C]
| | | |IncSubtensor{Inc;int64} [id BQ] │ │ └─ ···
| | | | |Elemwise{second,no_inplace} [id BR] │ ├─ Subtensor{::int64} [id BL] (outer_in_mit_mot-0)
| | | | | |Subtensor{int64::} [id BS] │ │ ├─ IncSubtensor{Inc;int64::} [id BM]
| | | | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0) │ │ │ ├─ Elemwise{second,no_inplace} [id BN]
| | | | | | |ScalarConstant{1} [id BT] │ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | | |InplaceDimShuffle{x,x} [id BU] │ │ │ │ │ └─ ···
| | | | | |TensorConstant{0.0} [id BV] │ │ │ │ └─ InplaceDimShuffle{x,x} [id BO]
| | | | |Elemwise{second} [id BW] │ │ │ │ └─ TensorConstant{0.0} [id BP]
| | | | | |Subtensor{int64} [id BX] │ │ │ ├─ IncSubtensor{Inc;int64} [id BQ]
| | | | | | |Subtensor{int64::} [id BS] │ │ │ │ ├─ Elemwise{second,no_inplace} [id BR]
| | | | | | |ScalarConstant{-1} [id BY] │ │ │ │ │ ├─ Subtensor{int64::} [id BS]
| | | | | |InplaceDimShuffle{x} [id BZ] │ │ │ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | | |Elemwise{second,no_inplace} [id CA] │ │ │ │ │ │ │ └─ ···
| | | | | |Sum{acc_dtype=float64} [id CB] │ │ │ │ │ │ └─ ScalarConstant{1} [id BT]
| | | | | | |Subtensor{int64} [id BX] │ │ │ │ │ └─ InplaceDimShuffle{x,x} [id BU]
| | | | | |TensorConstant{1.0} [id CC] │ │ │ │ │ └─ TensorConstant{0.0} [id BV]
| | | | |ScalarConstant{-1} [id BY] │ │ │ │ ├─ Elemwise{second} [id BW]
| | | |ScalarConstant{1} [id BT] │ │ │ │ │ ├─ Subtensor{int64} [id BX]
| | |ScalarConstant{-1} [id CD] │ │ │ │ │ │ ├─ Subtensor{int64::} [id BS]
| |Alloc [id CE] (outer_in_sit_sot-0) │ │ │ │ │ │ │ └─ ···
| | |TensorConstant{0.0} [id CF] │ │ │ │ │ │ └─ ScalarConstant{-1} [id BY]
| | |Elemwise{add,no_inplace} [id CG] │ │ │ │ │ └─ InplaceDimShuffle{x} [id BZ]
| | | |Elemwise{sub,no_inplace} [id C] │ │ │ │ │ └─ Elemwise{second,no_inplace} [id CA]
| | | |TensorConstant{1} [id CH] │ │ │ │ │ ├─ Sum{acc_dtype=float64} [id CB]
| | |Subtensor{int64} [id CI] │ │ │ │ │ │ └─ Subtensor{int64} [id BX]
| | |Shape [id CJ] │ │ │ │ │ │ └─ ···
| | | |A [id P] │ │ │ │ │ └─ TensorConstant{1.0} [id CC]
| | |ScalarConstant{0} [id CK] │ │ │ │ └─ ScalarConstant{-1} [id BY]
| |A [id P] (outer_in_non_seqs-0) │ │ │ └─ ScalarConstant{1} [id BT]
|ScalarConstant{-1} [id CL] │ │ └─ ScalarConstant{-1} [id CD]
│ ├─ Alloc [id CE] (outer_in_sit_sot-0)
│ │ ├─ TensorConstant{0.0} [id CF]
│ │ ├─ Elemwise{add,no_inplace} [id CG]
│ │ │ ├─ Elemwise{sub,no_inplace} [id C]
│ │ │ │ └─ ···
│ │ │ └─ TensorConstant{1} [id CH]
│ │ └─ Subtensor{int64} [id CI]
│ │ ├─ Shape [id CJ]
│ │ │ └─ A [id P]
│ │ └─ ScalarConstant{0} [id CK]
│ └─ A [id P] (outer_in_non_seqs-0)
└─ ScalarConstant{-1} [id CL]
Inner graphs: Inner graphs:
for{cpu,grad_of_scan_fn}.1 [id B] (outer_out_sit_sot-0) for{cpu,grad_of_scan_fn} [id B]
>Elemwise{add,no_inplace} [id CM] (inner_out_mit_mot-0-0) ← Elemwise{add,no_inplace} [id CM] (inner_out_mit_mot-0-0)
> |Elemwise{mul} [id CN] ├─ Elemwise{mul} [id CN]
> | |*2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0) │ ├─ *2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
> | |*5-<TensorType(float64, (?,))> [id CP] -> [id P] (inner_in_non_seqs-0) │ └─ *5-<TensorType(float64, (?,))> [id CP] -> [id P] (inner_in_non_seqs-0)
> |*3-<TensorType(float64, (?,))> [id CQ] -> [id BL] (inner_in_mit_mot-0-1) └─ *3-<TensorType(float64, (?,))> [id CQ] -> [id BL] (inner_in_mit_mot-0-1)
>Elemwise{add,no_inplace} [id CR] (inner_out_sit_sot-0) ← Elemwise{add,no_inplace} [id CR] (inner_out_sit_sot-0)
> |Elemwise{mul} [id CS] ├─ Elemwise{mul} [id CS]
> | |*2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0) │ ├─ *2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
> | |*0-<TensorType(float64, (?,))> [id CT] -> [id Z] (inner_in_seqs-0) │ └─ *0-<TensorType(float64, (?,))> [id CT] -> [id Z] (inner_in_seqs-0)
> |*4-<TensorType(float64, (?,))> [id CU] -> [id CE] (inner_in_sit_sot-0) └─ *4-<TensorType(float64, (?,))> [id CU] -> [id CE] (inner_in_sit_sot-0)
for{cpu,scan_fn} [id F] (outer_out_sit_sot-0) for{cpu,scan_fn} [id F]
>Elemwise{mul,no_inplace} [id CV] (inner_out_sit_sot-0) ← Elemwise{mul,no_inplace} [id CV] (inner_out_sit_sot-0)
> |*0-<TensorType(float64, (?,))> [id CT] -> [id H] (inner_in_sit_sot-0) ├─ *0-<TensorType(float64, (?,))> [id CT] -> [id H] (inner_in_sit_sot-0)
> |*1-<TensorType(float64, (?,))> [id CW] -> [id P] (inner_in_non_seqs-0)""" └─ *1-<TensorType(float64, (?,))> [id CW] -> [id P] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip() assert truth.strip() == out.strip()
...@@ -602,40 +642,39 @@ def test_debugprint_compiled_fn(): ...@@ -602,40 +642,39 @@ def test_debugprint_compiled_fn():
out = pytensor.function([M], out, updates=updates, mode="FAST_RUN") out = pytensor.function([M], out, updates=updates, mode="FAST_RUN")
expected_output = """forall_inplace,cpu,scan_fn} [id A] 2 (outer_out_sit_sot-0) expected_output = """forall_inplace,cpu,scan_fn} [id A] 2 (outer_out_sit_sot-0)
|TensorConstant{20000} [id B] (n_steps) ├─ TensorConstant{20000} [id B] (n_steps)
|TensorConstant{[ 0 ..998 19999]} [id C] (outer_in_seqs-0) ├─ TensorConstant{[ 0 ..998 19999]} [id C] (outer_in_seqs-0)
|IncSubtensor{InplaceSet;:int64:} [id D] 1 (outer_in_sit_sot-0) ├─ IncSubtensor{InplaceSet;:int64:} [id D] 1 (outer_in_sit_sot-0)
| |AllocEmpty{dtype='int64'} [id E] 0 │ ├─ AllocEmpty{dtype='int64'} [id E] 0
| | |TensorConstant{20000} [id B] │ │ └─ TensorConstant{20000} [id B]
| |TensorConstant{(1,) of 0} [id F] │ ├─ TensorConstant{(1,) of 0} [id F]
| |ScalarConstant{1} [id G] │ └─ ScalarConstant{1} [id G]
|<TensorType(float64, (20000, 2, 2))> [id H] (outer_in_non_seqs-0) └─ <TensorType(float64, (20000, 2, 2))> [id H] (outer_in_non_seqs-0)
Inner graphs: Inner graphs:
forall_inplace,cpu,scan_fn} [id A] (outer_out_sit_sot-0) forall_inplace,cpu,scan_fn} [id A]
>Elemwise{Composite} [id I] (inner_out_sit_sot-0) ← Elemwise{Composite} [id I] (inner_out_sit_sot-0)
> |TensorConstant{0} [id J] ├─ TensorConstant{0} [id J]
> |Subtensor{int64, int64, uint8} [id K] ├─ Subtensor{int64, int64, uint8} [id K]
> | |*2-<TensorType(float64, (20000, 2, 2))> [id L] -> [id H] (inner_in_non_seqs-0) │ ├─ *2-<TensorType(float64, (20000, 2, 2))> [id L] -> [id H] (inner_in_non_seqs-0)
> | |ScalarFromTensor [id M] │ ├─ ScalarFromTensor [id M]
> | | |*0-<TensorType(int64, ())> [id N] -> [id C] (inner_in_seqs-0) │ │ └─ *0-<TensorType(int64, ())> [id N] -> [id C] (inner_in_seqs-0)
> | |ScalarFromTensor [id O] │ ├─ ScalarFromTensor [id O]
> | | |*1-<TensorType(int64, ())> [id P] -> [id D] (inner_in_sit_sot-0) │ │ └─ *1-<TensorType(int64, ())> [id P] -> [id D] (inner_in_sit_sot-0)
> | |ScalarConstant{0} [id Q] │ └─ ScalarConstant{0} [id Q]
> |TensorConstant{1} [id R] └─ TensorConstant{1} [id R]
Elemwise{Composite} [id I] Elemwise{Composite} [id I]
>Switch [id S] ← Switch [id S]
> |LT [id T] ├─ LT [id T]
> | |<int64> [id U] │ ├─ <int64> [id U]
> | |<float64> [id V] │ └─ <float64> [id V]
> |<int64> [id W] ├─ <int64> [id W]
> |<int64> [id U] └─ <int64> [id U]
""" """
output_str = debugprint(out, file="str", print_op_info=True) output_str = debugprint(out, file="str", print_op_info=True)
print(output_str)
lines = output_str.split("\n") lines = output_str.split("\n")
for truth, out in zip(expected_output.split("\n"), lines): for truth, out in zip(expected_output.split("\n"), lines):
......
tests/test_printing.py
浏览文件 @ 1f2542eb
...@@ -145,12 +145,12 @@ def test_debugprint(): ...@@ -145,12 +145,12 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} [id 0] Elemwise{add,no_inplace} [id 0]
|Elemwise{add,no_inplace} [id 1] 'C' ├─ Elemwise{add,no_inplace} [id 1] 'C'
| |A [id 2] │ ├─ A [id 2]
| |B [id 3] │ └─ B [id 3]
|Elemwise{add,no_inplace} [id 4] └─ Elemwise{add,no_inplace} [id 4]
|D [id 5] ├─ D [id 5]
|E [id 6] └─ E [id 6]
""" """
).lstrip() ).lstrip()
...@@ -163,12 +163,12 @@ def test_debugprint(): ...@@ -163,12 +163,12 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} [id A] Elemwise{add,no_inplace} [id A]
|Elemwise{add,no_inplace} [id B] 'C' ├─ Elemwise{add,no_inplace} [id B] 'C'
| |A [id C] │ ├─ A [id C]
| |B [id D] │ └─ B [id D]
|Elemwise{add,no_inplace} [id E] └─ Elemwise{add,no_inplace} [id E]
|D [id F] ├─ D [id F]
|E [id G] └─ E [id G]
""" """
).lstrip() ).lstrip()
...@@ -181,10 +181,11 @@ def test_debugprint(): ...@@ -181,10 +181,11 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} [id A] Elemwise{add,no_inplace} [id A]
|Elemwise{add,no_inplace} [id B] 'C' ├─ Elemwise{add,no_inplace} [id B] 'C'
|Elemwise{add,no_inplace} [id C] │ └─ ···
|D [id D] └─ Elemwise{add,no_inplace} [id C]
|E [id E] ├─ D [id D]
└─ E [id E]
""" """
).lstrip() ).lstrip()
...@@ -196,12 +197,12 @@ def test_debugprint(): ...@@ -196,12 +197,12 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} Elemwise{add,no_inplace}
|Elemwise{add,no_inplace} 'C' ├─ Elemwise{add,no_inplace} 'C'
| |A │ ├─ A
| |B │ └─ B
|Elemwise{add,no_inplace} └─ Elemwise{add,no_inplace}
|D ├─ D
|E └─ E
""" """
).lstrip() ).lstrip()
...@@ -213,10 +214,10 @@ def test_debugprint(): ...@@ -213,10 +214,10 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} 0 [None] Elemwise{add,no_inplace} 0 [None]
|A [None] ├─ A [None]
|B [None] ├─ B [None]
|D [None] ├─ D [None]
|E [None] └─ E [None]
""" """
).lstrip() ).lstrip()
...@@ -231,10 +232,10 @@ def test_debugprint(): ...@@ -231,10 +232,10 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} 0 [None] Elemwise{add,no_inplace} 0 [None]
|A [None] ├─ A [None]
|B [None] ├─ B [None]
|D [None] ├─ D [None]
|E [None] └─ E [None]
""" """
).lstrip() ).lstrip()
...@@ -249,10 +250,10 @@ def test_debugprint(): ...@@ -249,10 +250,10 @@ def test_debugprint():
reference = dedent( reference = dedent(
r""" r"""
Elemwise{add,no_inplace} 0 [None] Elemwise{add,no_inplace} 0 [None]
|A [None] ├─ A [None]
|B [None] ├─ B [None]
|D [None] ├─ D [None]
|E [None] └─ E [None]
""" """
).lstrip() ).lstrip()
...@@ -274,26 +275,26 @@ def test_debugprint(): ...@@ -274,26 +275,26 @@ def test_debugprint():
exp_res = dedent( exp_res = dedent(
r""" r"""
Elemwise{Composite} 4 Elemwise{Composite} 4
|InplaceDimShuffle{x,0} v={0: [0]} 3 ├─ InplaceDimShuffle{x,0} v={0: [0]} 3
| |CGemv{inplace} d={0: [0]} 2 │ └─ CGemv{inplace} d={0: [0]} 2
| |AllocEmpty{dtype='float64'} 1 │ ├─ AllocEmpty{dtype='float64'} 1
| | |Shape_i{0} 0 │ │ └─ Shape_i{0} 0
| | |B │ │ └─ B
| |TensorConstant{1.0} │ ├─ TensorConstant{1.0}
| |B │ ├─ B
| |<TensorType(float64, (?,))> │ ├─ <TensorType(float64, (?,))>
| |TensorConstant{0.0} │ └─ TensorConstant{0.0}
|D ├─ D
|A └─ A
Inner graphs: Inner graphs:
Elemwise{Composite} Elemwise{Composite}
>add ← add
> |<float64> ├─ <float64>
> |sub └─ sub
> |<float64> ├─ <float64>
> |<float64> └─ <float64>
""" """
).lstrip() ).lstrip()
...@@ -314,10 +315,10 @@ def test_debugprint_id_type(): ...@@ -314,10 +315,10 @@ def test_debugprint_id_type():
s = s.getvalue() s = s.getvalue()
exp_res = f"""Elemwise{{add,no_inplace}} [id {e_at.auto_name}] exp_res = f"""Elemwise{{add,no_inplace}} [id {e_at.auto_name}]
|dot [id {d_at.auto_name}] ├─ dot [id {d_at.auto_name}]
| |<TensorType(float64, (?, ?))> [id {b_at.auto_name}] │ ├─ <TensorType(float64, (?, ?))> [id {b_at.auto_name}]
| |<TensorType(float64, (?,))> [id {a_at.auto_name}] │ └─ <TensorType(float64, (?,))> [id {a_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")] == [ assert [l.strip() for l in s.split("\n")] == [
...@@ -351,15 +352,15 @@ def test_debugprint_inner_graph(): ...@@ -351,15 +352,15 @@ def test_debugprint_inner_graph():
lines = output_str.split("\n") lines = output_str.split("\n")
exp_res = """MyInnerGraphOp [id A] exp_res = """MyInnerGraphOp [id A]
|3 [id B] ├─ 3 [id B]
|4 [id C] └─ 4 [id C]
Inner graphs: Inner graphs:
MyInnerGraphOp [id A] MyInnerGraphOp [id A]
>op2 [id D] 'igo1' ← op2 [id D] 'igo1'
> |*0-<MyType()> [id E] ├─ *0-<MyType()> [id E]
> |*1-<MyType()> [id F] └─ *1-<MyType()> [id F]
""" """
for exp_line, res_line in zip(exp_res.split("\n"), lines): for exp_line, res_line in zip(exp_res.split("\n"), lines):
...@@ -375,19 +376,19 @@ MyInnerGraphOp [id A] ...@@ -375,19 +376,19 @@ MyInnerGraphOp [id A]
lines = output_str.split("\n") lines = output_str.split("\n")
exp_res = """MyInnerGraphOp [id A] exp_res = """MyInnerGraphOp [id A]
|5 [id B] └─ 5 [id B]
Inner graphs: Inner graphs:
MyInnerGraphOp [id A] MyInnerGraphOp [id A]
>MyInnerGraphOp [id C] ← MyInnerGraphOp [id C]
> |*0-<MyType()> [id D] ├─ *0-<MyType()> [id D]
> |*1-<MyType()> [id E] └─ *1-<MyType()> [id E]
MyInnerGraphOp [id C] MyInnerGraphOp [id C]
>op2 [id F] 'igo1' ← op2 [id F] 'igo1'
> |*0-<MyType()> [id D] ├─ *0-<MyType()> [id D]
> |*1-<MyType()> [id E] └─ *1-<MyType()> [id E]
""" """
for exp_line, res_line in zip(exp_res.split("\n"), lines): for exp_line, res_line in zip(exp_res.split("\n"), lines):
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论