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.:
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
浏览文件 @ 1f2542eb
......@@ -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/scan/test_printing.py
浏览文件 @ 1f2542eb
......@@ -27,39 +27,42 @@ def test_debugprint_sitsot():
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A]
|Subtensor{int64::} [id B]
| |for{cpu,scan_fn} [id C] (outer_out_sit_sot-0)
| | |k [id D] (n_steps)
| | |IncSubtensor{Set;:int64:} [id E] (outer_in_sit_sot-0)
| | | |AllocEmpty{dtype='float64'} [id F]
| | | | |Elemwise{add,no_inplace} [id G]
| | | | | |k [id D]
| | | | | |Subtensor{int64} [id H]
| | | | | |Shape [id I]
| | | | | | |Unbroadcast{0} [id J]
| | | | | | |InplaceDimShuffle{x,0} [id K]
| | | | | | |Elemwise{second,no_inplace} [id L]
| | | | | | |A [id M]
| | | | | | |InplaceDimShuffle{x} [id N]
| | | | | | |TensorConstant{1.0} [id O]
| | | | | |ScalarConstant{0} [id P]
| | | | |Subtensor{int64} [id Q]
| | | | |Shape [id R]
| | | | | |Unbroadcast{0} [id J]
| | | | |ScalarConstant{1} [id S]
| | | |Unbroadcast{0} [id J]
| | | |ScalarFromTensor [id T]
| | | |Subtensor{int64} [id H]
| | |A [id M] (outer_in_non_seqs-0)
| |ScalarConstant{1} [id U]
|ScalarConstant{-1} [id V]
├─ Subtensor{int64::} [id B]
│ ├─ for{cpu,scan_fn} [id C] (outer_out_sit_sot-0)
│ │ ├─ k [id D] (n_steps)
│ │ ├─ IncSubtensor{Set;:int64:} [id E] (outer_in_sit_sot-0)
│ │ │ ├─ AllocEmpty{dtype='float64'} [id F]
│ │ │ │ ├─ Elemwise{add,no_inplace} [id G]
│ │ │ │ │ ├─ k [id D]
│ │ │ │ │ └─ Subtensor{int64} [id H]
│ │ │ │ │ ├─ Shape [id I]
│ │ │ │ │ │ └─ Unbroadcast{0} [id J]
│ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id K]
│ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id L]
│ │ │ │ │ │ ├─ A [id M]
│ │ │ │ │ │ └─ InplaceDimShuffle{x} [id N]
│ │ │ │ │ │ └─ TensorConstant{1.0} [id O]
│ │ │ │ │ └─ ScalarConstant{0} [id P]
│ │ │ │ └─ Subtensor{int64} [id Q]
│ │ │ │ ├─ Shape [id R]
│ │ │ │ │ └─ Unbroadcast{0} [id J]
│ │ │ │ │ └─ ···
│ │ │ │ └─ ScalarConstant{1} [id S]
│ │ │ ├─ Unbroadcast{0} [id J]
│ │ │ │ └─ ···
│ │ │ └─ ScalarFromTensor [id T]
│ │ │ └─ Subtensor{int64} [id H]
│ │ │ └─ ···
│ │ └─ A [id M] (outer_in_non_seqs-0)
│ └─ ScalarConstant{1} [id U]
└─ ScalarConstant{-1} [id V]
Inner graphs:
for{cpu,scan_fn} [id C] (outer_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)
> |*1-<TensorType(float64, (?,))> [id Y] -> [id M] (inner_in_non_seqs-0)"""
for{cpu,scan_fn} [id C]
← Elemwise{mul,no_inplace} [id W] (inner_out_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)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -82,39 +85,42 @@ def test_debugprint_sitsot_no_extra_info():
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A]
|Subtensor{int64::} [id B]
| |for{cpu,scan_fn} [id C]
| | |k [id D]
| | |IncSubtensor{Set;:int64:} [id E]
| | | |AllocEmpty{dtype='float64'} [id F]
| | | | |Elemwise{add,no_inplace} [id G]
| | | | | |k [id D]
| | | | | |Subtensor{int64} [id H]
| | | | | |Shape [id I]
| | | | | | |Unbroadcast{0} [id J]
| | | | | | |InplaceDimShuffle{x,0} [id K]
| | | | | | |Elemwise{second,no_inplace} [id L]
| | | | | | |A [id M]
| | | | | | |InplaceDimShuffle{x} [id N]
| | | | | | |TensorConstant{1.0} [id O]
| | | | | |ScalarConstant{0} [id P]
| | | | |Subtensor{int64} [id Q]
| | | | |Shape [id R]
| | | | | |Unbroadcast{0} [id J]
| | | | |ScalarConstant{1} [id S]
| | | |Unbroadcast{0} [id J]
| | | |ScalarFromTensor [id T]
| | | |Subtensor{int64} [id H]
| | |A [id M]
| |ScalarConstant{1} [id U]
|ScalarConstant{-1} [id V]
├─ Subtensor{int64::} [id B]
│ ├─ for{cpu,scan_fn} [id C]
│ │ ├─ k [id D]
│ │ ├─ IncSubtensor{Set;:int64:} [id E]
│ │ │ ├─ AllocEmpty{dtype='float64'} [id F]
│ │ │ │ ├─ Elemwise{add,no_inplace} [id G]
│ │ │ │ │ ├─ k [id D]
│ │ │ │ │ └─ Subtensor{int64} [id H]
│ │ │ │ │ ├─ Shape [id I]
│ │ │ │ │ │ └─ Unbroadcast{0} [id J]
│ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id K]
│ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id L]
│ │ │ │ │ │ ├─ A [id M]
│ │ │ │ │ │ └─ InplaceDimShuffle{x} [id N]
│ │ │ │ │ │ └─ TensorConstant{1.0} [id O]
│ │ │ │ │ └─ ScalarConstant{0} [id P]
│ │ │ │ └─ Subtensor{int64} [id Q]
│ │ │ │ ├─ Shape [id R]
│ │ │ │ │ └─ Unbroadcast{0} [id J]
│ │ │ │ │ └─ ···
│ │ │ │ └─ ScalarConstant{1} [id S]
│ │ │ ├─ Unbroadcast{0} [id J]
│ │ │ │ └─ ···
│ │ │ └─ ScalarFromTensor [id T]
│ │ │ └─ Subtensor{int64} [id H]
│ │ │ └─ ···
│ │ └─ A [id M]
│ └─ ScalarConstant{1} [id U]
└─ ScalarConstant{-1} [id V]
Inner graphs:
for{cpu,scan_fn} [id C]
>Elemwise{mul,no_inplace} [id W]
> |*0-<TensorType(float64, (?,))> [id X] -> [id E]
> |*1-<TensorType(float64, (?,))> [id Y] -> [id M]"""
← Elemwise{mul,no_inplace} [id W]
├─ *0-<TensorType(float64, (?,))> [id X] -> [id E]
└─ *1-<TensorType(float64, (?,))> [id Y] -> [id M]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -142,42 +148,47 @@ def test_debugprint_nitsot():
lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A]
|for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
|Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
| |Subtensor{int64} [id D]
| | |Shape [id E]
| | | |Subtensor{int64::} [id F] 'coefficients[0:]'
| | | |coefficients [id G]
| | | |ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I]
| |Subtensor{int64} [id J]
| |Shape [id K]
| | |Subtensor{int64::} [id L]
| | |ARange{dtype='int64'} [id M]
| | | |TensorConstant{0} [id N]
| | | |TensorConstant{10000} [id O]
| | | |TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] (outer_in_seqs-0)
| |Subtensor{int64::} [id F] 'coefficients[0:]'
| |ScalarFromTensor [id T]
| |Elemwise{scalar_minimum,no_inplace} [id C]
|Subtensor{:int64:} [id U] (outer_in_seqs-1)
| |Subtensor{int64::} [id L]
| |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)
└─ for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
├─ Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
│ ├─ Subtensor{int64} [id D]
│ │ ├─ Shape [id E]
│ │ │ └─ Subtensor{int64::} [id F] 'coefficients[0:]'
│ │ │ ├─ coefficients [id G]
│ │ │ └─ ScalarConstant{0} [id H]
│ │ └─ ScalarConstant{0} [id I]
│ └─ Subtensor{int64} [id J]
│ ├─ Shape [id K]
│ │ └─ Subtensor{int64::} [id L]
│ │ ├─ ARange{dtype='int64'} [id M]
│ │ │ ├─ TensorConstant{0} [id N]
│ │ │ ├─ TensorConstant{10000} [id O]
│ │ │ └─ TensorConstant{1} [id P]
│ │ └─ ScalarConstant{0} [id Q]
│ └─ ScalarConstant{0} [id R]
├─ Subtensor{:int64:} [id S] (outer_in_seqs-0)
│ ├─ Subtensor{int64::} [id F] 'coefficients[0:]'
│ │ └─ ···
│ └─ ScalarFromTensor [id T]
│ └─ Elemwise{scalar_minimum,no_inplace} [id C]
│ └─ ···
├─ Subtensor{:int64:} [id U] (outer_in_seqs-1)
│ ├─ Subtensor{int64::} [id L]
│ │ └─ ···
│ └─ 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:
for{cpu,scan_fn} [id B] (outer_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)
> |Elemwise{pow,no_inplace} [id Z]
> |*2-<TensorType(float64, ())> [id BA] -> [id W] (inner_in_non_seqs-0)
> |*1-<TensorType(int64, ())> [id BB] -> [id U] (inner_in_seqs-1)"""
for{cpu,scan_fn} [id B]
← Elemwise{mul,no_inplace} [id X] (inner_out_nit_sot-0)
├─ *0-<TensorType(float64, ())> [id Y] -> [id S] (inner_in_seqs-0)
└─ Elemwise{pow,no_inplace} [id Z]
├─ *2-<TensorType(float64, ())> [id BA] -> [id W] (inner_in_non_seqs-0)
└─ *1-<TensorType(int64, ())> [id BB] -> [id U] (inner_in_seqs-1)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -215,76 +226,84 @@ def test_debugprint_nested_scans():
lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A]
|for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
|Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
| |Subtensor{int64} [id D]
| | |Shape [id E]
| | | |Subtensor{int64::} [id F] 'c[0:]'
| | | |c [id G]
| | | |ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I]
| |Subtensor{int64} [id J]
| |Shape [id K]
| | |Subtensor{int64::} [id L]
| | |ARange{dtype='int64'} [id M]
| | | |TensorConstant{0} [id N]
| | | |TensorConstant{10} [id O]
| | | |TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] (outer_in_seqs-0)
| |Subtensor{int64::} [id F] 'c[0:]'
| |ScalarFromTensor [id T]
| |Elemwise{scalar_minimum,no_inplace} [id C]
|Subtensor{:int64:} [id U] (outer_in_seqs-1)
| |Subtensor{int64::} [id L]
| |ScalarFromTensor [id V]
| |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)
└─ for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
├─ Elemwise{scalar_minimum,no_inplace} [id C] (outer_in_nit_sot-0)
│ ├─ Subtensor{int64} [id D]
│ │ ├─ Shape [id E]
│ │ │ └─ Subtensor{int64::} [id F] 'c[0:]'
│ │ │ ├─ c [id G]
│ │ │ └─ ScalarConstant{0} [id H]
│ │ └─ ScalarConstant{0} [id I]
│ └─ Subtensor{int64} [id J]
│ ├─ Shape [id K]
│ │ └─ Subtensor{int64::} [id L]
│ │ ├─ ARange{dtype='int64'} [id M]
│ │ │ ├─ TensorConstant{0} [id N]
│ │ │ ├─ TensorConstant{10} [id O]
│ │ │ └─ TensorConstant{1} [id P]
│ │ └─ ScalarConstant{0} [id Q]
│ └─ ScalarConstant{0} [id R]
├─ Subtensor{:int64:} [id S] (outer_in_seqs-0)
│ ├─ Subtensor{int64::} [id F] 'c[0:]'
│ │ └─ ···
│ └─ ScalarFromTensor [id T]
│ └─ Elemwise{scalar_minimum,no_inplace} [id C]
│ └─ ···
├─ Subtensor{:int64:} [id U] (outer_in_seqs-1)
│ ├─ Subtensor{int64::} [id L]
│ │ └─ ···
│ └─ ScalarFromTensor [id V]
│ └─ 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:
for{cpu,scan_fn} [id B] (outer_out_nit_sot-0)
>Elemwise{mul,no_inplace} [id Y] (inner_out_nit_sot-0)
> |InplaceDimShuffle{x} [id Z]
> | |*0-<TensorType(float64, ())> [id BA] -> [id S] (inner_in_seqs-0)
> |Elemwise{pow,no_inplace} [id BB]
> |Subtensor{int64} [id BC]
> | |Subtensor{int64::} [id BD]
> | | |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)
> | | | |IncSubtensor{Set;:int64:} [id BG] (outer_in_sit_sot-0)
> | | | | |AllocEmpty{dtype='float64'} [id BH]
> | | | | | |Elemwise{add,no_inplace} [id BI]
> | | | | | | |*3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1)
> | | | | | | |Subtensor{int64} [id BJ]
> | | | | | | |Shape [id BK]
> | | | | | | | |Unbroadcast{0} [id BL]
> | | | | | | | |InplaceDimShuffle{x,0} [id BM]
> | | | | | | | |Elemwise{second,no_inplace} [id BN]
> | | | | | | | |*2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0)
> | | | | | | | |InplaceDimShuffle{x} [id BP]
> | | | | | | | |TensorConstant{1.0} [id BQ]
> | | | | | | |ScalarConstant{0} [id BR]
> | | | | | |Subtensor{int64} [id BS]
> | | | | | |Shape [id BT]
> | | | | | | |Unbroadcast{0} [id BL]
> | | | | | |ScalarConstant{1} [id BU]
> | | | | |Unbroadcast{0} [id BL]
> | | | | |ScalarFromTensor [id BV]
> | | | | |Subtensor{int64} [id BJ]
> | | | |*2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
> | | |ScalarConstant{1} [id BW]
> | |ScalarConstant{-1} [id BX]
> |InplaceDimShuffle{x} [id BY]
> |*1-<TensorType(int64, ())> [id BZ] -> [id U] (inner_in_seqs-1)
for{cpu,scan_fn} [id BE] (outer_out_sit_sot-0)
>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{cpu,scan_fn} [id B]
← Elemwise{mul,no_inplace} [id Y] (inner_out_nit_sot-0)
├─ InplaceDimShuffle{x} [id Z]
│ └─ *0-<TensorType(float64, ())> [id BA] -> [id S] (inner_in_seqs-0)
└─ Elemwise{pow,no_inplace} [id BB]
├─ Subtensor{int64} [id BC]
│ ├─ Subtensor{int64::} [id BD]
│ │ ├─ 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)
│ │ │ ├─ IncSubtensor{Set;:int64:} [id BG] (outer_in_sit_sot-0)
│ │ │ │ ├─ AllocEmpty{dtype='float64'} [id BH]
│ │ │ │ │ ├─ Elemwise{add,no_inplace} [id BI]
│ │ │ │ │ │ ├─ *3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1)
│ │ │ │ │ │ └─ Subtensor{int64} [id BJ]
│ │ │ │ │ │ ├─ Shape [id BK]
│ │ │ │ │ │ │ └─ Unbroadcast{0} [id BL]
│ │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id BM]
│ │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id BN]
│ │ │ │ │ │ │ ├─ *2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0)
│ │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id BP]
│ │ │ │ │ │ │ └─ TensorConstant{1.0} [id BQ]
│ │ │ │ │ │ └─ ScalarConstant{0} [id BR]
│ │ │ │ │ └─ Subtensor{int64} [id BS]
│ │ │ │ │ ├─ Shape [id BT]
│ │ │ │ │ │ └─ Unbroadcast{0} [id BL]
│ │ │ │ │ │ └─ ···
│ │ │ │ │ └─ ScalarConstant{1} [id BU]
│ │ │ │ ├─ Unbroadcast{0} [id BL]
│ │ │ │ │ └─ ···
│ │ │ │ └─ ScalarFromTensor [id BV]
│ │ │ │ └─ Subtensor{int64} [id BJ]
│ │ │ │ └─ ···
│ │ │ └─ *2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
│ │ └─ ScalarConstant{1} [id BW]
│ └─ ScalarConstant{-1} [id BX]
└─ InplaceDimShuffle{x} [id BY]
└─ *1-<TensorType(int64, ())> [id BZ] -> [id U] (inner_in_seqs-1)
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):
assert truth.strip() == out.strip()
......@@ -296,86 +315,94 @@ def test_debugprint_nested_scans():
)
lines = output_str.split("\n")
expected_output = """-c [id A]
-k [id B]
-A [id C]
expected_output = """→ c [id A]
→ k [id B]
→ A [id C]
Sum{acc_dtype=float64} [id D] 13
|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)
| |Subtensor{int64} [id G] 6
| | |Shape [id H] 5
| | | |Subtensor{int64::} [id I] 'c[0:]' 4
| | | |c [id A]
| | | |ScalarConstant{0} [id J]
| | |ScalarConstant{0} [id K]
| |Subtensor{int64} [id L] 3
| |Shape [id M] 2
| | |Subtensor{int64::} [id N] 1
| | |ARange{dtype='int64'} [id O] 0
| | | |TensorConstant{0} [id P]
| | | |TensorConstant{10} [id Q]
| | | |TensorConstant{1} [id R]
| | |ScalarConstant{0} [id S]
| |ScalarConstant{0} [id T]
|Subtensor{:int64:} [id U] 11 (outer_in_seqs-0)
| |Subtensor{int64::} [id I] 'c[0:]' 4
| |ScalarFromTensor [id V] 10
| |Elemwise{scalar_minimum,no_inplace} [id F] 7
|Subtensor{:int64:} [id W] 9 (outer_in_seqs-1)
| |Subtensor{int64::} [id N] 1
| |ScalarFromTensor [id X] 8
| |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)
└─ 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)
│ ├─ Subtensor{int64} [id G] 6
│ │ ├─ Shape [id H] 5
│ │ │ └─ Subtensor{int64::} [id I] 'c[0:]' 4
│ │ │ ├─ c [id A]
│ │ │ └─ ScalarConstant{0} [id J]
│ │ └─ ScalarConstant{0} [id K]
│ └─ Subtensor{int64} [id L] 3
│ ├─ Shape [id M] 2
│ │ └─ Subtensor{int64::} [id N] 1
│ │ ├─ ARange{dtype='int64'} [id O] 0
│ │ │ ├─ TensorConstant{0} [id P]
│ │ │ ├─ TensorConstant{10} [id Q]
│ │ │ └─ TensorConstant{1} [id R]
│ │ └─ ScalarConstant{0} [id S]
│ └─ ScalarConstant{0} [id T]
├─ Subtensor{:int64:} [id U] 11 (outer_in_seqs-0)
│ ├─ Subtensor{int64::} [id I] 'c[0:]' 4
│ │ └─ ···
│ └─ ScalarFromTensor [id V] 10
│ └─ Elemwise{scalar_minimum,no_inplace} [id F] 7
│ └─ ···
├─ Subtensor{:int64:} [id W] 9 (outer_in_seqs-1)
│ ├─ Subtensor{int64::} [id N] 1
│ │ └─ ···
│ └─ ScalarFromTensor [id X] 8
│ └─ 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:
for{cpu,scan_fn} [id E] (outer_out_nit_sot-0)
-*0-<TensorType(float64, ())> [id Y] -> [id U] (inner_in_seqs-0)
-*1-<TensorType(int64, ())> [id Z] -> [id W] (inner_in_seqs-1)
-*2-<TensorType(float64, (?,))> [id BA] -> [id C] (inner_in_non_seqs-0)
-*3-<TensorType(int32, ())> [id BB] -> [id B] (inner_in_non_seqs-1)
>Elemwise{mul,no_inplace} [id BC] (inner_out_nit_sot-0)
> |InplaceDimShuffle{x} [id BD]
> | |*0-<TensorType(float64, ())> [id Y] (inner_in_seqs-0)
> |Elemwise{pow,no_inplace} [id BE]
> |Subtensor{int64} [id BF]
> | |Subtensor{int64::} [id BG]
> | | |for{cpu,scan_fn} [id BH] (outer_out_sit_sot-0)
> | | | |*3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1) (n_steps)
> | | | |IncSubtensor{Set;:int64:} [id BI] (outer_in_sit_sot-0)
> | | | | |AllocEmpty{dtype='float64'} [id BJ]
> | | | | | |Elemwise{add,no_inplace} [id BK]
> | | | | | | |*3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1)
> | | | | | | |Subtensor{int64} [id BL]
> | | | | | | |Shape [id BM]
> | | | | | | | |Unbroadcast{0} [id BN]
> | | | | | | | |InplaceDimShuffle{x,0} [id BO]
> | | | | | | | |Elemwise{second,no_inplace} [id BP]
> | | | | | | | |*2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0)
> | | | | | | | |InplaceDimShuffle{x} [id BQ]
> | | | | | | | |TensorConstant{1.0} [id BR]
> | | | | | | |ScalarConstant{0} [id BS]
> | | | | | |Subtensor{int64} [id BT]
> | | | | | |Shape [id BU]
> | | | | | | |Unbroadcast{0} [id BN]
> | | | | | |ScalarConstant{1} [id BV]
> | | | | |Unbroadcast{0} [id BN]
> | | | | |ScalarFromTensor [id BW]
> | | | | |Subtensor{int64} [id BL]
> | | | |*2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
> | | |ScalarConstant{1} [id BX]
> | |ScalarConstant{-1} [id BY]
> |InplaceDimShuffle{x} [id BZ]
> |*1-<TensorType(int64, ())> [id Z] (inner_in_seqs-1)
for{cpu,scan_fn} [id BH] (outer_out_sit_sot-0)
-*0-<TensorType(float64, (?,))> [id CA] -> [id BI] (inner_in_sit_sot-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{cpu,scan_fn} [id E]
→ *0-<TensorType(float64, ())> [id Y] -> [id U] (inner_in_seqs-0)
→ *1-<TensorType(int64, ())> [id Z] -> [id W] (inner_in_seqs-1)
→ *2-<TensorType(float64, (?,))> [id BA] -> [id C] (inner_in_non_seqs-0)
→ *3-<TensorType(int32, ())> [id BB] -> [id B] (inner_in_non_seqs-1)
← Elemwise{mul,no_inplace} [id BC] (inner_out_nit_sot-0)
├─ InplaceDimShuffle{x} [id BD]
│ └─ *0-<TensorType(float64, ())> [id Y] (inner_in_seqs-0)
└─ Elemwise{pow,no_inplace} [id BE]
├─ Subtensor{int64} [id BF]
│ ├─ Subtensor{int64::} [id BG]
│ │ ├─ for{cpu,scan_fn} [id BH] (outer_out_sit_sot-0)
│ │ │ ├─ *3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1) (n_steps)
│ │ │ ├─ IncSubtensor{Set;:int64:} [id BI] (outer_in_sit_sot-0)
│ │ │ │ ├─ AllocEmpty{dtype='float64'} [id BJ]
│ │ │ │ │ ├─ Elemwise{add,no_inplace} [id BK]
│ │ │ │ │ │ ├─ *3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1)
│ │ │ │ │ │ └─ Subtensor{int64} [id BL]
│ │ │ │ │ │ ├─ Shape [id BM]
│ │ │ │ │ │ │ └─ Unbroadcast{0} [id BN]
│ │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id BO]
│ │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id BP]
│ │ │ │ │ │ │ ├─ *2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0)
│ │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id BQ]
│ │ │ │ │ │ │ └─ TensorConstant{1.0} [id BR]
│ │ │ │ │ │ └─ ScalarConstant{0} [id BS]
│ │ │ │ │ └─ Subtensor{int64} [id BT]
│ │ │ │ │ ├─ Shape [id BU]
│ │ │ │ │ │ └─ Unbroadcast{0} [id BN]
│ │ │ │ │ │ └─ ···
│ │ │ │ │ └─ ScalarConstant{1} [id BV]
│ │ │ │ ├─ Unbroadcast{0} [id BN]
│ │ │ │ │ └─ ···
│ │ │ │ └─ ScalarFromTensor [id BW]
│ │ │ │ └─ Subtensor{int64} [id BL]
│ │ │ │ └─ ···
│ │ │ └─ *2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
│ │ └─ ScalarConstant{1} [id BX]
│ └─ ScalarConstant{-1} [id BY]
└─ InplaceDimShuffle{x} [id BZ]
└─ *1-<TensorType(int64, ())> [id Z] (inner_in_seqs-1)
for{cpu,scan_fn} [id BH]
→ *0-<TensorType(float64, (?,))> [id CA] -> [id BI] (inner_in_sit_sot-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):
assert truth.strip() == out.strip()
......@@ -403,53 +430,54 @@ def test_debugprint_mitsot():
lines = output_str.split("\n")
expected_output = """Elemwise{add,no_inplace} [id A]
|Subtensor{int64::} [id B]
| |for{cpu,scan_fn}.0 [id C] (outer_out_mit_sot-0)
| | |TensorConstant{5} [id D] (n_steps)
| | |IncSubtensor{Set;:int64:} [id E] (outer_in_mit_sot-0)
| | | |AllocEmpty{dtype='int64'} [id F]
| | | | |Elemwise{add,no_inplace} [id G]
| | | | |TensorConstant{5} [id D]
| | | | |Subtensor{int64} [id H]
| | | | |Shape [id I]
| | | | | |Subtensor{:int64:} [id J]
| | | | | |<TensorType(int64, (?,))> [id K]
| | | | | |ScalarConstant{2} [id L]
| | | | |ScalarConstant{0} [id M]
| | | |Subtensor{:int64:} [id J]
| | | |ScalarFromTensor [id N]
| | | |Subtensor{int64} [id H]
| | |IncSubtensor{Set;:int64:} [id O] (outer_in_mit_sot-1)
| | |AllocEmpty{dtype='int64'} [id P]
| | | |Elemwise{add,no_inplace} [id Q]
| | | |TensorConstant{5} [id D]
| | | |Subtensor{int64} [id R]
| | | |Shape [id S]
| | | | |Subtensor{:int64:} [id T]
| | | | |<TensorType(int64, (?,))> [id U]
| | | | |ScalarConstant{2} [id V]
| | | |ScalarConstant{0} [id W]
| | |Subtensor{:int64:} [id T]
| | |ScalarFromTensor [id X]
| | |Subtensor{int64} [id R]
| |ScalarConstant{2} [id Y]
|Subtensor{int64::} [id Z]
|for{cpu,scan_fn}.1 [id C] (outer_out_mit_sot-1)
|ScalarConstant{2} [id BA]
├─ Subtensor{int64::} [id B]
│ ├─ for{cpu,scan_fn}.0 [id C] (outer_out_mit_sot-0)
│ │ ├─ TensorConstant{5} [id D] (n_steps)
│ │ ├─ IncSubtensor{Set;:int64:} [id E] (outer_in_mit_sot-0)
│ │ │ ├─ AllocEmpty{dtype='int64'} [id F]
│ │ │ │ └─ Elemwise{add,no_inplace} [id G]
│ │ │ │ ├─ TensorConstant{5} [id D]
│ │ │ │ └─ Subtensor{int64} [id H]
│ │ │ │ ├─ Shape [id I]
│ │ │ │ │ └─ Subtensor{:int64:} [id J]
│ │ │ │ │ ├─ <TensorType(int64, (?,))> [id K]
│ │ │ │ │ └─ ScalarConstant{2} [id L]
│ │ │ │ └─ ScalarConstant{0} [id M]
│ │ │ ├─ Subtensor{:int64:} [id J]
│ │ │ │ └─ ···
│ │ │ └─ ScalarFromTensor [id N]
│ │ │ └─ Subtensor{int64} [id H]
│ │ │ └─ ···
│ │ └─ IncSubtensor{Set;:int64:} [id O] (outer_in_mit_sot-1)
│ │ ├─ AllocEmpty{dtype='int64'} [id P]
│ │ │ └─ Elemwise{add,no_inplace} [id Q]
│ │ │ ├─ TensorConstant{5} [id D]
│ │ │ └─ Subtensor{int64} [id R]
│ │ │ ├─ Shape [id S]
│ │ │ │ └─ Subtensor{:int64:} [id T]
│ │ │ │ ├─ <TensorType(int64, (?,))> [id U]
│ │ │ │ └─ ScalarConstant{2} [id V]
│ │ │ └─ ScalarConstant{0} [id W]
│ │ ├─ Subtensor{:int64:} [id T]
│ │ │ └─ ···
│ │ └─ ScalarFromTensor [id X]
│ │ └─ Subtensor{int64} [id R]
│ │ └─ ···
│ └─ 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:
for{cpu,scan_fn}.0 [id C] (outer_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)
> |*0-<TensorType(int64, ())> [id BD] -> [id E] (inner_in_mit_sot-0-0)
>Elemwise{add,no_inplace} [id BE] (inner_out_mit_sot-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)
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{cpu,scan_fn} [id C]
← Elemwise{add,no_inplace} [id BB] (inner_out_mit_sot-0)
├─ *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)
← Elemwise{add,no_inplace} [id BE] (inner_out_mit_sot-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)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -474,106 +502,118 @@ def test_debugprint_mitmot():
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A]
|for{cpu,grad_of_scan_fn}.1 [id B] (outer_out_sit_sot-0)
| |Elemwise{sub,no_inplace} [id C] (n_steps)
| | |Subtensor{int64} [id D]
| | | |Shape [id E]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | |k [id G] (n_steps)
| | | | |IncSubtensor{Set;:int64:} [id H] (outer_in_sit_sot-0)
| | | | | |AllocEmpty{dtype='float64'} [id I]
| | | | | | |Elemwise{add,no_inplace} [id J]
| | | | | | | |k [id G]
| | | | | | | |Subtensor{int64} [id K]
| | | | | | | |Shape [id L]
| | | | | | | | |Unbroadcast{0} [id M]
| | | | | | | | |InplaceDimShuffle{x,0} [id N]
| | | | | | | | |Elemwise{second,no_inplace} [id O]
| | | | | | | | |A [id P]
| | | | | | | | |InplaceDimShuffle{x} [id Q]
| | | | | | | | |TensorConstant{1.0} [id R]
| | | | | | | |ScalarConstant{0} [id S]
| | | | | | |Subtensor{int64} [id T]
| | | | | | |Shape [id U]
| | | | | | | |Unbroadcast{0} [id M]
| | | | | | |ScalarConstant{1} [id V]
| | | | | |Unbroadcast{0} [id M]
| | | | | |ScalarFromTensor [id W]
| | | | | |Subtensor{int64} [id K]
| | | | |A [id P] (outer_in_non_seqs-0)
| | | |ScalarConstant{0} [id X]
| | |TensorConstant{1} [id Y]
| |Subtensor{:int64:} [id Z] (outer_in_seqs-0)
| | |Subtensor{::int64} [id BA]
| | | |Subtensor{:int64:} [id BB]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | |ScalarConstant{-1} [id BC]
| | | |ScalarConstant{-1} [id BD]
| | |ScalarFromTensor [id BE]
| | |Elemwise{sub,no_inplace} [id C]
| |Subtensor{:int64:} [id BF] (outer_in_seqs-1)
| | |Subtensor{:int64:} [id BG]
| | | |Subtensor{::int64} [id BH]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | |ScalarConstant{-1} [id BI]
| | | |ScalarConstant{-1} [id BJ]
| | |ScalarFromTensor [id BK]
| | |Elemwise{sub,no_inplace} [id C]
| |Subtensor{::int64} [id BL] (outer_in_mit_mot-0)
| | |IncSubtensor{Inc;int64::} [id BM]
| | | |Elemwise{second,no_inplace} [id BN]
| | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | |InplaceDimShuffle{x,x} [id BO]
| | | | |TensorConstant{0.0} [id BP]
| | | |IncSubtensor{Inc;int64} [id BQ]
| | | | |Elemwise{second,no_inplace} [id BR]
| | | | | |Subtensor{int64::} [id BS]
| | | | | | |for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
| | | | | | |ScalarConstant{1} [id BT]
| | | | | |InplaceDimShuffle{x,x} [id BU]
| | | | | |TensorConstant{0.0} [id BV]
| | | | |Elemwise{second} [id BW]
| | | | | |Subtensor{int64} [id BX]
| | | | | | |Subtensor{int64::} [id BS]
| | | | | | |ScalarConstant{-1} [id BY]
| | | | | |InplaceDimShuffle{x} [id BZ]
| | | | | |Elemwise{second,no_inplace} [id CA]
| | | | | |Sum{acc_dtype=float64} [id CB]
| | | | | | |Subtensor{int64} [id BX]
| | | | | |TensorConstant{1.0} [id CC]
| | | | |ScalarConstant{-1} [id BY]
| | | |ScalarConstant{1} [id BT]
| | |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]
├─ for{cpu,grad_of_scan_fn}.1 [id B] (outer_out_sit_sot-0)
│ ├─ Elemwise{sub,no_inplace} [id C] (n_steps)
│ │ ├─ Subtensor{int64} [id D]
│ │ │ ├─ Shape [id E]
│ │ │ │ └─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
│ │ │ │ ├─ k [id G] (n_steps)
│ │ │ │ ├─ IncSubtensor{Set;:int64:} [id H] (outer_in_sit_sot-0)
│ │ │ │ │ ├─ AllocEmpty{dtype='float64'} [id I]
│ │ │ │ │ │ ├─ Elemwise{add,no_inplace} [id J]
│ │ │ │ │ │ │ ├─ k [id G]
│ │ │ │ │ │ │ └─ Subtensor{int64} [id K]
│ │ │ │ │ │ │ ├─ Shape [id L]
│ │ │ │ │ │ │ │ └─ Unbroadcast{0} [id M]
│ │ │ │ │ │ │ │ └─ InplaceDimShuffle{x,0} [id N]
│ │ │ │ │ │ │ │ └─ Elemwise{second,no_inplace} [id O]
│ │ │ │ │ │ │ │ ├─ A [id P]
│ │ │ │ │ │ │ │ └─ InplaceDimShuffle{x} [id Q]
│ │ │ │ │ │ │ │ └─ TensorConstant{1.0} [id R]
│ │ │ │ │ │ │ └─ ScalarConstant{0} [id S]
│ │ │ │ │ │ └─ Subtensor{int64} [id T]
│ │ │ │ │ │ ├─ Shape [id U]
│ │ │ │ │ │ │ └─ Unbroadcast{0} [id M]
│ │ │ │ │ │ │ └─ ···
│ │ │ │ │ │ └─ ScalarConstant{1} [id V]
│ │ │ │ │ ├─ Unbroadcast{0} [id M]
│ │ │ │ │ │ └─ ···
│ │ │ │ │ └─ ScalarFromTensor [id W]
│ │ │ │ │ └─ Subtensor{int64} [id K]
│ │ │ │ │ └─ ···
│ │ │ │ └─ A [id P] (outer_in_non_seqs-0)
│ │ │ └─ ScalarConstant{0} [id X]
│ │ └─ TensorConstant{1} [id Y]
│ ├─ Subtensor{:int64:} [id Z] (outer_in_seqs-0)
│ │ ├─ Subtensor{::int64} [id BA]
│ │ │ ├─ Subtensor{:int64:} [id BB]
│ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
│ │ │ │ │ └─ ···
│ │ │ │ └─ ScalarConstant{-1} [id BC]
│ │ │ └─ ScalarConstant{-1} [id BD]
│ │ └─ ScalarFromTensor [id BE]
│ │ └─ Elemwise{sub,no_inplace} [id C]
│ │ └─ ···
│ ├─ Subtensor{:int64:} [id BF] (outer_in_seqs-1)
│ │ ├─ Subtensor{:int64:} [id BG]
│ │ │ ├─ Subtensor{::int64} [id BH]
│ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
│ │ │ │ │ └─ ···
│ │ │ │ └─ ScalarConstant{-1} [id BI]
│ │ │ └─ ScalarConstant{-1} [id BJ]
│ │ └─ ScalarFromTensor [id BK]
│ │ └─ Elemwise{sub,no_inplace} [id C]
│ │ └─ ···
│ ├─ Subtensor{::int64} [id BL] (outer_in_mit_mot-0)
│ │ ├─ IncSubtensor{Inc;int64::} [id BM]
│ │ │ ├─ Elemwise{second,no_inplace} [id BN]
│ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
│ │ │ │ │ └─ ···
│ │ │ │ └─ InplaceDimShuffle{x,x} [id BO]
│ │ │ │ └─ TensorConstant{0.0} [id BP]
│ │ │ ├─ IncSubtensor{Inc;int64} [id BQ]
│ │ │ │ ├─ Elemwise{second,no_inplace} [id BR]
│ │ │ │ │ ├─ Subtensor{int64::} [id BS]
│ │ │ │ │ │ ├─ for{cpu,scan_fn} [id F] (outer_out_sit_sot-0)
│ │ │ │ │ │ │ └─ ···
│ │ │ │ │ │ └─ ScalarConstant{1} [id BT]
│ │ │ │ │ └─ InplaceDimShuffle{x,x} [id BU]
│ │ │ │ │ └─ TensorConstant{0.0} [id BV]
│ │ │ │ ├─ Elemwise{second} [id BW]
│ │ │ │ │ ├─ Subtensor{int64} [id BX]
│ │ │ │ │ │ ├─ Subtensor{int64::} [id BS]
│ │ │ │ │ │ │ └─ ···
│ │ │ │ │ │ └─ ScalarConstant{-1} [id BY]
│ │ │ │ │ └─ InplaceDimShuffle{x} [id BZ]
│ │ │ │ │ └─ Elemwise{second,no_inplace} [id CA]
│ │ │ │ │ ├─ Sum{acc_dtype=float64} [id CB]
│ │ │ │ │ │ └─ Subtensor{int64} [id BX]
│ │ │ │ │ │ └─ ···
│ │ │ │ │ └─ TensorConstant{1.0} [id CC]
│ │ │ │ └─ ScalarConstant{-1} [id BY]
│ │ │ └─ ScalarConstant{1} [id BT]
│ │ └─ 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:
for{cpu,grad_of_scan_fn}.1 [id B] (outer_out_sit_sot-0)
>Elemwise{add,no_inplace} [id CM] (inner_out_mit_mot-0-0)
> |Elemwise{mul} [id CN]
> | |*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)
> |*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{mul} [id CS]
> | |*2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
> | |*0-<TensorType(float64, (?,))> [id CT] -> [id Z] (inner_in_seqs-0)
> |*4-<TensorType(float64, (?,))> [id CU] -> [id CE] (inner_in_sit_sot-0)
for{cpu,scan_fn} [id F] (outer_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)
> |*1-<TensorType(float64, (?,))> [id CW] -> [id P] (inner_in_non_seqs-0)"""
for{cpu,grad_of_scan_fn} [id B]
← Elemwise{add,no_inplace} [id CM] (inner_out_mit_mot-0-0)
├─ Elemwise{mul} [id CN]
│ ├─ *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)
└─ *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{mul} [id CS]
│ ├─ *2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
│ └─ *0-<TensorType(float64, (?,))> [id CT] -> [id Z] (inner_in_seqs-0)
└─ *4-<TensorType(float64, (?,))> [id CU] -> [id CE] (inner_in_sit_sot-0)
for{cpu,scan_fn} [id F]
← Elemwise{mul,no_inplace} [id CV] (inner_out_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)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -602,40 +642,39 @@ def test_debugprint_compiled_fn():
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)
|TensorConstant{20000} [id B] (n_steps)
|TensorConstant{[ 0 ..998 19999]} [id C] (outer_in_seqs-0)
|IncSubtensor{InplaceSet;:int64:} [id D] 1 (outer_in_sit_sot-0)
| |AllocEmpty{dtype='int64'} [id E] 0
| | |TensorConstant{20000} [id B]
| |TensorConstant{(1,) of 0} [id F]
| |ScalarConstant{1} [id G]
|<TensorType(float64, (20000, 2, 2))> [id H] (outer_in_non_seqs-0)
├─ TensorConstant{20000} [id B] (n_steps)
├─ TensorConstant{[ 0 ..998 19999]} [id C] (outer_in_seqs-0)
├─ IncSubtensor{InplaceSet;:int64:} [id D] 1 (outer_in_sit_sot-0)
│ ├─ AllocEmpty{dtype='int64'} [id E] 0
│ │ └─ TensorConstant{20000} [id B]
│ ├─ TensorConstant{(1,) of 0} [id F]
│ └─ ScalarConstant{1} [id G]
└─ <TensorType(float64, (20000, 2, 2))> [id H] (outer_in_non_seqs-0)
Inner graphs:
forall_inplace,cpu,scan_fn} [id A] (outer_out_sit_sot-0)
>Elemwise{Composite} [id I] (inner_out_sit_sot-0)
> |TensorConstant{0} [id J]
> |Subtensor{int64, int64, uint8} [id K]
> | |*2-<TensorType(float64, (20000, 2, 2))> [id L] -> [id H] (inner_in_non_seqs-0)
> | |ScalarFromTensor [id M]
> | | |*0-<TensorType(int64, ())> [id N] -> [id C] (inner_in_seqs-0)
> | |ScalarFromTensor [id O]
> | | |*1-<TensorType(int64, ())> [id P] -> [id D] (inner_in_sit_sot-0)
> | |ScalarConstant{0} [id Q]
> |TensorConstant{1} [id R]
forall_inplace,cpu,scan_fn} [id A]
← Elemwise{Composite} [id I] (inner_out_sit_sot-0)
├─ TensorConstant{0} [id J]
├─ Subtensor{int64, int64, uint8} [id K]
│ ├─ *2-<TensorType(float64, (20000, 2, 2))> [id L] -> [id H] (inner_in_non_seqs-0)
│ ├─ ScalarFromTensor [id M]
│ │ └─ *0-<TensorType(int64, ())> [id N] -> [id C] (inner_in_seqs-0)
│ ├─ ScalarFromTensor [id O]
│ │ └─ *1-<TensorType(int64, ())> [id P] -> [id D] (inner_in_sit_sot-0)
│ └─ ScalarConstant{0} [id Q]
└─ TensorConstant{1} [id R]
Elemwise{Composite} [id I]
>Switch [id S]
> |LT [id T]
> | |<int64> [id U]
> | |<float64> [id V]
> |<int64> [id W]
> |<int64> [id U]
← Switch [id S]
├─ LT [id T]
│ ├─ <int64> [id U]
│ └─ <float64> [id V]
├─ <int64> [id W]
└─ <int64> [id U]
"""
output_str = debugprint(out, file="str", print_op_info=True)
print(output_str)
lines = output_str.split("\n")
for truth, out in zip(expected_output.split("\n"), lines):
......
tests/test_printing.py
浏览文件 @ 1f2542eb
......@@ -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):
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论