diff --git a/src/frontends/tensorflow_common/src/op/rank.cpp b/src/frontends/tensorflow_common/src/op/rank.cpp index 1c37dfa5a0e878..1adbcd43d2fc04 100644 --- a/src/frontends/tensorflow_common/src/op/rank.cpp +++ b/src/frontends/tensorflow_common/src/op/rank.cpp @@ -3,8 +3,11 @@ // #include "common_op_table.hpp" +#include "helper_ops/complex_type_mark.hpp" +#include "openvino/op/constant.hpp" #include "openvino/op/shape_of.hpp" #include "openvino/op/squeeze.hpp" +#include "openvino/op/subtract.hpp" using namespace std; using namespace ov::op; @@ -15,8 +18,21 @@ namespace tensorflow { namespace op { ov::OutputVector translate_rank_op(const NodeContext& node) { - default_op_checks(node, 1, {"Rank"}); + default_op_checks(node, 1, {"Rank"}, true); auto input = node.get_input(0); + auto complex_type_mark = as_type_ptr(input.get_node_shared_ptr()); + if (complex_type_mark) { + input = complex_type_mark->input_value(0); + auto input_shape = make_shared(input, ov::element::i32); + + auto unsqueeze_input_rank = make_shared(input_shape, ov::element::i32); + auto input_rank_with_complex = make_shared(unsqueeze_input_rank); + // eliminate the extra dimension + auto input_rank = + make_shared(input_rank_with_complex, make_shared(ov::element::i32, Shape{}, 1)); + set_node_name(node.get_name(), input_rank); + return {input_rank->output(0)}; + } auto input_shape = make_shared(input, ov::element::i32); auto unsqueeze_input_rank = make_shared(input_shape, ov::element::i32); auto input_rank = make_shared(unsqueeze_input_rank); diff --git a/tests/layer_tests/tensorflow_tests/test_tf_Rank.py b/tests/layer_tests/tensorflow_tests/test_tf_Rank.py index e047b8ecb4ed8b..0a6ef2b08f6f57 100644 --- a/tests/layer_tests/tensorflow_tests/test_tf_Rank.py +++ b/tests/layer_tests/tensorflow_tests/test_tf_Rank.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +import numpy as np import tensorflow as tf from common.tf_layer_test_class import CommonTFLayerTest @@ -33,3 +34,44 @@ def test_rank_basic(self, params, ie_device, precision, ir_version, temp_dir, self._test(*self.create_rank_net(**params), ie_device, precision, ir_version, temp_dir=temp_dir, use_legacy_frontend=use_legacy_frontend) + +class TestComplexRank(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_rank_net(self, input_shape): + tf.compat.v1.reset_default_graph() + # Create the graph and model + with tf.compat.v1.Session() as sess: + input_real = tf.compat.v1.placeholder(tf.float32, input_shape, 'param_real') + input_imag = tf.compat.v1.placeholder(tf.float32, input_shape, 'param_imag') + input = tf.raw_ops.Complex(real=input_real, imag=input_imag) + tf.raw_ops.Rank(input=input) + 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=[3, 4, 5, 6]) + ] + + @pytest.mark.parametrize("params", test_data_basic) + @pytest.mark.precommit_tf_fe + @pytest.mark.nightly + def test_complex_rank(self, params, ie_device, precision, ir_version, temp_dir, + use_legacy_frontend): + self._test(*self.create_rank_net(**params), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_legacy_frontend=use_legacy_frontend)