diff --git a/src/frontends/tensorflow/src/transformations/switch_merge_resolve.cpp b/src/frontends/tensorflow/src/transformations/switch_merge_resolve.cpp index 35c2cd1b7f2..07e5ac31ee5 100644 --- a/src/frontends/tensorflow/src/transformations/switch_merge_resolve.cpp +++ b/src/frontends/tensorflow/src/transformations/switch_merge_resolve.cpp @@ -55,12 +55,14 @@ void generate_if_clusters(const shared_ptr& ov_model, continue; } auto eliminated_markers = merge_node->get_eliminated_cond_flow_marker(); - if (eliminated_markers.size() != 1) { - continue; - } - auto eliminated_marker = eliminated_markers[0]; - auto switch_nodes = merge_node->get_switch_nodes_set_by_cond_index(eliminated_marker); + // combine all Switch nodes for which conditional flow is resolved + // by the current Merge node + SetOfSwitchNodes switch_nodes; + for (const auto& eliminated_marker : eliminated_markers) { + auto curr_switch_nodes = merge_node->get_switch_nodes_set_by_cond_index(eliminated_marker); + switch_nodes.insert(curr_switch_nodes.begin(), curr_switch_nodes.end()); + } // insert into clusters ClusterType combined_cluster = {switch_nodes, {merge_node}}; @@ -121,11 +123,13 @@ void insert_result_before_merge(const shared_ptr& merge_node, "[TensorFlow Frontend] internal error: Merge node " + merge_node_name + " does not have conditional flow marker"); - // get eliminated marker and check that it is the single one + // get eliminated marker and check that eliminated marker exists + // Merge node may contain several eliminated markers, in this case it means some Switch nodes have different + // condition nodes and values generated by this condition nodes are identical auto merge_cf_marker = get_cf_marker(merge_node); - FRONT_END_GENERAL_CHECK(merge_cf_marker.merge_eliminated_markers.size() == 1, + FRONT_END_GENERAL_CHECK(merge_cf_marker.merge_eliminated_markers.size() > 0, "[TensorFlow Frontend] internal error: Merge node " + merge_node_name + - " does not contain the single eliminated marker"); + " does not contain any eliminated marker"); auto eliminated_marker = merge_cf_marker.merge_eliminated_markers.begin()->first; // check that producer contains the same conditional flow marker diff --git a/tests/layer_tests/tensorflow_tests/test_tf_SwitchMerge.py b/tests/layer_tests/tensorflow_tests/test_tf_SwitchMerge.py new file mode 100644 index 00000000000..8acc25c3a60 --- /dev/null +++ b/tests/layer_tests/tensorflow_tests/test_tf_SwitchMerge.py @@ -0,0 +1,65 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +import tensorflow as tf +from common.tf_layer_test_class import CommonTFLayerTest + + +class TestSwitchMerge(CommonTFLayerTest): + def _prepare_input(self, inputs_info): + assert 'x' in inputs_info + x_shape = inputs_info['x'] + inputs_data = {} + rng = np.random.default_rng() + inputs_data['x'] = rng.integers(-10, 10, x_shape).astype(self.x_type) + return inputs_data + + def merge_eliminating_several_cond_flows_net(self, x_shape, x_type, cond_value): + self.x_type = x_type + tf.compat.v1.reset_default_graph() + # Create the graph and model + with tf.compat.v1.Session() as sess: + x = tf.compat.v1.placeholder(x_type, x_shape, 'x') + cond = tf.constant(cond_value, dtype=tf.bool) + switch_false, switch_true = tf.raw_ops.Switch(data=x, pred=cond) + + cond2 = tf.constant(cond_value, dtype=tf.bool) + switch2_false, switch2_true = tf.raw_ops.Switch(data=cond2, pred=cond2) + with tf.control_dependencies([switch2_true]): + const_sub = tf.constant(5, dtype=x_type) + with tf.control_dependencies([switch2_false]): + const_add = tf.constant(2, dtype=x_type) + + add = tf.raw_ops.AddV2(x=switch_false, y=const_add) + sub = tf.raw_ops.Sub(x=switch_true, y=const_sub) + merge = tf.raw_ops.Merge(inputs=[add, sub]) + const_main = tf.constant(1, dtype=x_type) + tf.raw_ops.AddV2(x=merge[0], y=const_main) + tf.compat.v1.global_variables_initializer() + tf_net = sess.graph_def + + return tf_net, None + + test_data_basic = [ + dict(x_shape=[]), + dict(x_shape=[2]), + dict(x_shape=[4, 3]), + ] + + @pytest.mark.parametrize("cond_value", [ + True, False + ]) + @pytest.mark.parametrize("x_type", [ + np.float32, np.int32 + ]) + @pytest.mark.parametrize("params", test_data_basic) + @pytest.mark.precommit_tf_fe + @pytest.mark.nightly + def test_merge_eliminating_several_cond_flows(self, params, cond_value, x_type, ie_device, precision, ir_version, + temp_dir, + use_new_frontend, use_old_api): + self._test(*self.merge_eliminating_several_cond_flows_net(**params, cond_value=cond_value, x_type=x_type), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_new_frontend=use_new_frontend, use_old_api=use_old_api)