def test_backward_bfs_multi_consumer_data_nodes(self):
# Placeholder-> Mul -> Result
# Const -/ \- Result2
graph = build_graph(
{
**regular_op_with_shaped_data('parameter', [1], {
'op': 'Parameter'
}),
**valued_const_with_data('const', int64_array([5])),
**regular_op_with_shaped_data('mul', [1], {'op': 'Mul'}),
**result('result'),
**result('result2'),
}, [
*connect('parameter', '0:mul'),
*connect('const', '1:mul'),
*connect('mul:0', 'result'),
*connect_data('mul', 'result2'),
])
res = common_bfs(Node(graph, 'result'), ['Mul'], ['Parameter'],
is_backward=True,
attr_to_check='op',
follow_multi_consumer_data_nodes=True)
self.assertTrue(
len(res) == 1,
'The multi-consumer data node "mul_d" was not followed')
res = common_bfs(Node(graph, 'result'), ['Mul'], ['Parameter'],
is_backward=True,
attr_to_check='op')
self.assertTrue(
len(res) == 0, 'The multi-consumer data node "mul_d" was followed')
def test_backward_bfs_check_op_instead_of_type(self):
# Placeholder->ScaleShift->Mul1->Add1---->Concat
# `----------->Add2->Mul2--'
graph = build_graph(nodes_attributes,
[('placeholder_1', 'placeholder_1_data'),
('placeholder_1_data', 'add_2'),
('scaleshift_1_w', 'scaleshift_1'),
('scaleshift_1', 'scaleshift_1_data'),
('scaleshift_1_data', 'mul_1'),
('mul_1', 'mul_1_data'), ('mul_1_data', 'add_1'),
('add_1', 'add_1_data'), ('add_2', 'add_2_data'),
('add_2_data', 'mul_2'), ('mul_2', 'mul_2_data'),
('add_1_data', 'concat_1'),
('mul_2_data', 'concat_1'),
('concat_1', 'concat_1_data'),
('concat_1_data', 'op_output')])
res = common_bfs(Node(graph, 'concat_1'), ['Mul', 'Add'],
['Parameter'],
is_backward=True,
attr_to_check='op')
self.assertTrue(len(res) == 0, 'Smth went wrong with bfs')
res = common_bfs(Node(graph, 'concat_1'), ['Mul'], ['Add'],
is_backward=True,
attr_to_check='op')
self.assertTrue(
len(res) == 2
and all([res[x].id in ['add_1', 'add_2']
for x in range(len(res))]),
'Add operations was not found by bfs')
res = common_bfs(Node(graph, 'concat_1'), ['ScaleShift'], ['Add'],
is_backward=True,
attr_to_check='op')
self.assertTrue(len(res) == 0, 'BFS shouldn\'t find any operations')
res = common_bfs(Node(graph, 'concat_1'), [], ['Add'],
allowed_all=True,
is_backward=True,
attr_to_check='op')
self.assertTrue(
len(res) == 2
and all([res[x].id in ['add_1', 'add_2']
for x in range(len(res))]),
'Add operations was not found by bfs')
res = common_bfs(Node(graph, 'concat_1'), ['ScaleShift'],
['ScaleShift'],
is_backward=True,
attr_to_check='op')
self.assertTrue(
len(res) == 0,
'No one node should be found! But bfs found {} nodes'.format(
len(res)))
def parameter_unchanged_after_iteration(loop_node: Node, body_parameter: Node):
"""
Checks if the body Parameter node is connected to some body Result and the data provided to Result is not
changed between iterations. The data is considered unchanged if:
1. There is no back edge for this Parameter OR
2. There is a back edge from some Result to Parameter and there are only Identity ops in between or
Parameter is connected to Result directly.
:param loop_node: the Loop node to check
:param body_parameter: the body Parameter node
:return: the result of the check
"""
assert body_parameter.id in loop_node.body
assert body_parameter.soft_get('op') == 'Parameter'
if not any([attr['to_layer'] == body_parameter.soft_get('internal_layer_id') for attr in loop_node.back_edges]):
return True
for back_edge_attrs in loop_node.back_edges:
if back_edge_attrs['to_layer'] == body_parameter.soft_get('internal_layer_id'):
result_internal_id = back_edge_attrs['from_layer']
result_nodes = loop_node.body.get_op_nodes(internal_layer_id=result_internal_id)
assert len(result_nodes) == 1, 'There should be exactly one node with id {}, but there are {}' \
''.format(result_internal_id, len(result_nodes))
result_node = result_nodes[0]
# check that the Result node consumes data from Parameter node directly or through Identity operations
parameters = common_bfs(result_node, ['Identity'], ['Parameter'], is_backward=True, attr_to_check='op',
follow_multi_consumer_data_nodes=True)
if any([node.soft_get('internal_layer_id') == body_parameter.internal_layer_id for node in parameters]):
return True
return False