From 10e5936deece6b1e5bb140fd801c7549be81b46e Mon Sep 17 00:00:00 2001
From: Yegor Kruglov <yegor.kruglov@intel.com>
Date: Thu, 24 Jun 2021 18:53:13 +0300
Subject: [PATCH] [MO] Fixes in Merge infer function (#6125)

* move output shape setting to the conditions

* e2e fix

* updated unittests
---
 model-optimizer/extensions/ops/merge.py       |  2 +-
 .../unit_tests/extensions/ops/merge_test.py   | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/model-optimizer/extensions/ops/merge.py b/model-optimizer/extensions/ops/merge.py
index 53134bf8a96..e0a992dc4cf 100644
--- a/model-optimizer/extensions/ops/merge.py
+++ b/model-optimizer/extensions/ops/merge.py
@@ -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
diff --git a/model-optimizer/unit_tests/extensions/ops/merge_test.py b/model-optimizer/unit_tests/extensions/ops/merge_test.py
index 8a395cc178b..73c33a57402 100644
--- a/model-optimizer/unit_tests/extensions/ops/merge_test.py
+++ b/model-optimizer/unit_tests/extensions/ops/merge_test.py
@@ -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)