[MO] Fixes in Merge infer function (#6125)
* move output shape setting to the conditions * e2e fix * updated unittests
This commit is contained in:
1 parent
7211cd3aa6
commit
10e5936dee
2 files changed
+30
-1
No files matched your search
@@ -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)
|
||||
Reference in new issue
Block a user