[MO] Fixes in Merge infer function (#6125)

* move output shape setting to the conditions

* e2e fix

* updated unittests
This commit is contained in:
Yegor Kruglov authored and GitHub committed 2021-06-24 18:53:13 +03:00
1 parent 7211cd3aa6
commit 10e5936dee
2 files changed
+30 -1

No files matched your search

+1 -1
View File
@@ -24,6 +24,7 @@ class Merge(Op):
# we infer only through executable input nodes
inferred_nodes = [n for n in node.in_nodes().values() if n['is_partial_inferred']]
assert len(inferred_nodes) != 0
tensor = inferred_nodes[0]
if len(inferred_nodes) < len(node.in_nodes()):
node['is_not_fully_inferred'] = True
@@ -38,7 +39,6 @@ class Merge(Op):
if all([np.all(tensor.value == n.value) for n in inferred_and_executable]):
node.out_node().value = tensor.value.copy() if tensor.has_valid('value') else None
tensor = inferred_nodes[0]
node.out_node().shape = int64_array(tensor.shape)
@staticmethod
@@ -6,6 +6,7 @@ import unittest
import numpy as np
from extensions.ops.merge import Merge
from mo.front.common.partial_infer.utils import int64_array
from mo.graph.graph import Node
from mo.utils.ir_engine.compare_graphs import compare_graphs
from unit_tests.utils.graph import build_graph_with_attrs
@@ -86,3 +87,31 @@ class TestMerge(unittest.TestCase):
('merge', {'is_not_fully_inferred': False})])
(flag, resp) = compare_graphs(graph, graph_ref, 'merge_output', check_op_attrs=True)
self.assertTrue(flag, resp)
def test_merge_infer_only_second_executable(self):
graph = build_graph_with_attrs(
nodes_with_attrs=self.nodes,
edges_with_attrs=self.edges,
update_nodes_attributes=[
('first', {'executable': False, 'value': np.ones([2, 2]), 'shape': int64_array([2, 2])}),
('second', {'executable': True, 'value': np.zeros([4, 4]), 'shape': int64_array([4, 4])})
]
)
ref_graph = build_graph_with_attrs(
nodes_with_attrs=self.nodes,
edges_with_attrs=self.edges,
update_nodes_attributes=[
('first', {'executable': False, 'value': np.ones([2, 2]), 'shape': int64_array([2, 2])}),
('second', {'executable': True, 'value': np.zeros([4, 4]), 'shape': int64_array([4, 4])}),
('merge', {'is_not_fully_inferred': False}),
('merge_output', {'shape': int64_array([4, 4]), 'value': np.zeros([4, 4])})
]
)
tested_class = Merge(graph=graph, attrs={})
node = Node(graph, 'merge')
tested_class.merge_infer(node)
(flag, resp) = compare_graphs(graph, ref_graph, 'merge_output', check_op_attrs=True)
self.assertTrue(flag, resp)