From 81e236f51ed25b11127b11decfda7c4cb9d2b5ee Mon Sep 17 00:00:00 2001 From: Luca Tam <155265433+LucaTamSapienza@users.noreply.github.com> Date: Tue, 12 Mar 2024 22:49:24 +0100 Subject: [PATCH] [TF FE] Supported complex tensors for Size operations (#23357) ### Details: - Added support for complex tensors and tests for it. ### Tickets: - #23244 --------- Co-authored-by: Roman Kazantsev --- .../tensorflow_common/src/op/size.cpp | 22 +++++++++- .../tensorflow_tests/test_tf_Size.py | 42 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/src/frontends/tensorflow_common/src/op/size.cpp b/src/frontends/tensorflow_common/src/op/size.cpp index 4a56f4d5aba086..00abf883a70ce9 100644 --- a/src/frontends/tensorflow_common/src/op/size.cpp +++ b/src/frontends/tensorflow_common/src/op/size.cpp @@ -3,7 +3,9 @@ // #include "common_op_table.hpp" +#include "helper_ops/complex_type_mark.hpp" #include "openvino/op/constant.hpp" +#include "openvino/op/divide.hpp" #include "openvino/op/reduce_prod.hpp" #include "openvino/op/shape_of.hpp" #include "openvino/op/unsqueeze.hpp" @@ -19,12 +21,29 @@ namespace op { ov::OutputVector translate_size_op(const NodeContext& node) { // Size operation computes a number of elements in the input tensor - default_op_checks(node, 1, {"Size"}); + default_op_checks(node, 1, {"Size"}, true); auto input = node.get_input(0); + auto complex_type_mark = as_type_ptr(input.get_node_shared_ptr()); + // retrive attribute of the output type auto out_type = node.get_attribute("out_type", element::i32); + if (complex_type_mark) { + input = complex_type_mark->input_value(0); + + // compute the input tensor size + auto shape_of = make_shared(input, out_type); + auto axis = make_shared(element::i32, Shape{}, 0); + auto complex_size = make_shared(shape_of, axis); + + // need to divide the size by 2 because real and imaginary parts are counted separately + auto complex_size_divided_by_two = + make_shared(complex_size, make_shared(element::i32, Shape{}, 2)); + set_node_name(node.get_name(), complex_size_divided_by_two); + return {complex_size_divided_by_two->output(0)}; + } + // introduce extra dimension in order to compute size in case of a scalar input auto const_zero = make_shared(element::i32, Shape{1}, 0); input = make_shared(input, const_zero); @@ -33,6 +52,7 @@ ov::OutputVector translate_size_op(const NodeContext& node) { auto shape_of = make_shared(input, out_type); auto axis = make_shared(element::i32, Shape{}, 0); auto size = make_shared(shape_of, axis); + set_node_name(node.get_name(), size); return {size}; } diff --git a/tests/layer_tests/tensorflow_tests/test_tf_Size.py b/tests/layer_tests/tensorflow_tests/test_tf_Size.py index 0622625e78a6ba..19705556f04176 100644 --- a/tests/layer_tests/tensorflow_tests/test_tf_Size.py +++ b/tests/layer_tests/tensorflow_tests/test_tf_Size.py @@ -43,3 +43,45 @@ def test_size_basic(self, params, ie_device, precision, ir_version, temp_dir, self._test(*self.create_size_net(**params), ie_device, precision, ir_version, temp_dir=temp_dir, use_legacy_frontend=use_legacy_frontend) + +class TestComplexSize(CommonTFLayerTest): + def _prepare_input(self, inputs_info): + rng = np.random.default_rng() + assert 'param_real:0' in inputs_info + assert 'param_imag:0' in inputs_info + param_real_shape_1 = inputs_info['param_real:0'] + param_imag_shape_1 = inputs_info['param_imag:0'] + inputs_data = {} + inputs_data['param_real:0'] = 4 * rng.random(param_real_shape_1).astype(np.float32) - 2 + inputs_data['param_imag:0'] = 4 * rng.random(param_imag_shape_1).astype(np.float32) - 2 + return inputs_data + + def create_complex_size_net(self, input_shape): + tf.compat.v1.reset_default_graph() + # Create the graph and model + with tf.compat.v1.Session() as sess: + param_real = tf.compat.v1.placeholder(np.float32, input_shape, 'param_real') + param_imag = tf.compat.v1.placeholder(np.float32, input_shape, 'param_imag') + complex = tf.raw_ops.Complex(real=param_real, imag=param_imag) + tf.raw_ops.Size(input=complex) + tf.compat.v1.global_variables_initializer() + tf_net = sess.graph_def + + return tf_net, None + + test_data_basic = [ + dict(input_shape=[]), + dict(input_shape=[1]), + dict(input_shape=[2, 6]), + dict(input_shape=[2, 4, 5]), + ] + + @pytest.mark.parametrize("params", test_data_basic) + @pytest.mark.precommit_tf_fe + @pytest.mark.nightly + def test_complex_size(self, params, ie_device, precision, ir_version, temp_dir, + use_legacy_frontend): + self._test( + *self.create_complex_size_net(**params), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_legacy_frontend=use_legacy_frontend)