From d264eb40b8184d0e6722e5adc3f33943f4c7db4d Mon Sep 17 00:00:00 2001 From: Adithya Hegde Kota Date: Sat, 4 May 2024 20:55:08 +0530 Subject: [PATCH] [TF FE]: Support complex tensors for SegmentSum operations (#23676) ### Details: - *Edited Tests for Segment Sum* Resolves #23241 --------- Co-authored-by: Roman Kazantsev --- .../tensorflow_common/src/op/segment_sum.cpp | 14 ++++- .../tensorflow_tests/test_tf_SegmentSum.py | 54 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/src/frontends/tensorflow_common/src/op/segment_sum.cpp b/src/frontends/tensorflow_common/src/op/segment_sum.cpp index 6abe729cc1e..457528ba4ce 100644 --- a/src/frontends/tensorflow_common/src/op/segment_sum.cpp +++ b/src/frontends/tensorflow_common/src/op/segment_sum.cpp @@ -3,6 +3,7 @@ // #include "common_op_table.hpp" +#include "helper_ops/complex_type_mark.hpp" #include "openvino/op/add.hpp" #include "openvino/op/constant.hpp" #include "openvino/op/embedding_segments_sum.hpp" @@ -21,7 +22,7 @@ namespace frontend { namespace tensorflow { namespace op { OutputVector translate_segment_sum_op(const NodeContext& node) { - default_op_checks(node, 2, {"SegmentSum"}); + default_op_checks(node, 2, {"SegmentSum"}, true); auto data = node.get_input(0); auto segment_ids = node.get_input(1); @@ -45,6 +46,17 @@ OutputVector translate_segment_sum_op(const NodeContext& node) { auto num_indices = make_shared(segment_ids_shape, squeeze_axis); auto indices = make_shared(const_zero, num_indices, const_one, indices_type); + auto complex_type_mark = as_type_ptr(data.get_node_shared_ptr()); + if (complex_type_mark) { + element::Type complex_part_type = complex_type_mark->get_complex_part_type(); + data = complex_type_mark->input_value(0); + auto emb_segment_sum_complex = make_shared(data, indices, segment_ids, num_segments); + auto emb_segment_sum_complex_output = + make_shared(emb_segment_sum_complex->output(0), complex_part_type); + set_node_name(node.get_name(), emb_segment_sum_complex); + return {emb_segment_sum_complex_output}; + } + auto emb_segment_sum = make_shared(data, indices, segment_ids, num_segments); set_node_name(node.get_name(), emb_segment_sum); return {emb_segment_sum}; diff --git a/tests/layer_tests/tensorflow_tests/test_tf_SegmentSum.py b/tests/layer_tests/tensorflow_tests/test_tf_SegmentSum.py index 2304c89e0f8..2c06b1a45ca 100644 --- a/tests/layer_tests/tensorflow_tests/test_tf_SegmentSum.py +++ b/tests/layer_tests/tensorflow_tests/test_tf_SegmentSum.py @@ -69,3 +69,57 @@ class TestSegmentSum(CommonTFLayerTest): self._test(*self.create_segment_sum_net(**params), ie_device, precision, ir_version, temp_dir=temp_dir, use_legacy_frontend=use_legacy_frontend) + + +class TestSegmentSumComplex(CommonTFLayerTest): + def _prepare_input(self, inputs_info): + rng = np.random.default_rng() + assert "data_real:0" in inputs_info + assert "data_imag:0" in inputs_info + assert "segment_ids:0" in inputs_info + data_real_shape = inputs_info["data_real:0"] + data_imag_shape = inputs_info["data_imag:0"] + segment_ids_shape = inputs_info["segment_ids:0"] + inputs_data = {} + inputs_data["data_real:0"] = rng.random(data_real_shape).astype(np.float32) + inputs_data["data_imag:0"] = rng.random(data_imag_shape).astype(np.float32) + inputs_data["segment_ids:0"] = np.sort(rng.integers( + low=0, high=data_real_shape[0], size=segment_ids_shape + ).astype(np.int32)) + return inputs_data + + def create_segment_sum_net(self, data_shape, segment_ids_shape): + tf.compat.v1.reset_default_graph() + # Create the graph and model + with tf.compat.v1.Session() as sess: + data_real = tf.compat.v1.placeholder(tf.float32,data_shape,'data_real') + data_imag = tf.compat.v1.placeholder(tf.float32,data_shape,'data_imag') + data = tf.raw_ops.Complex(real=data_real,imag=data_imag) + segment_ids = tf.compat.v1.placeholder(tf.int32, segment_ids_shape, "segment_ids") + segment_sum_output = tf.raw_ops.SegmentSum(data=data, segment_ids=segment_ids) + tf.raw_ops.Real(input=segment_sum_output) + tf.raw_ops.Imag(input=segment_sum_output) + tf.compat.v1.global_variables_initializer().run() + tf_net = sess.graph_def + + return tf_net, None + + test_data_basic = [ + dict(data_shape=[2, 3], segment_ids_shape=[2]), + dict(data_shape=[3, 2], segment_ids_shape=[3]), + dict(data_shape=[3, 1, 2], segment_ids_shape=[3]), + dict(data_shape=[4, 2, 1], segment_ids_shape=[4]), + ] + + @pytest.mark.parametrize("params", test_data_basic) + @pytest.mark.precommit + @pytest.mark.nightly + def test_complex_segment_sum(self, params, ie_device, precision, ir_version, temp_dir, + use_legacy_frontend): + if use_legacy_frontend: + pytest.skip("SegmentSum operation is not supported via legacy frontend.") + if ie_device == 'GPU': + pytest.skip("GPU error: to_shape was called on a dynamic shape") + self._test(*self.create_segment_sum_net(**params), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_legacy_frontend=use_legacy_frontend) \ No newline at end of file