Skip to content

Commit

Permalink
[TF FE] Add complex type support for rank op (#23118)
Browse files Browse the repository at this point in the history
### Details:
 - Add complex type support for rank op
 - Add corresponding test of it

### Tickets:
 - I think I don't have a ticket id

### Related issues

#22956

---------

Co-authored-by: Roman Kazantsev <[email protected]>
  • Loading branch information
AsakusaRinne and rkazants authored Mar 1, 2024
1 parent 4edb040 commit 7d35d2a
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 1 deletion.
18 changes: 17 additions & 1 deletion src/frontends/tensorflow_common/src/op/rank.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<ComplexTypeMark>(input.get_node_shared_ptr());
if (complex_type_mark) {
input = complex_type_mark->input_value(0);
auto input_shape = make_shared<v3::ShapeOf>(input, ov::element::i32);

auto unsqueeze_input_rank = make_shared<v3::ShapeOf>(input_shape, ov::element::i32);
auto input_rank_with_complex = make_shared<v0::Squeeze>(unsqueeze_input_rank);
// eliminate the extra dimension
auto input_rank =
make_shared<v1::Subtract>(input_rank_with_complex, make_shared<v0::Constant>(ov::element::i32, Shape{}, 1));
set_node_name(node.get_name(), input_rank);
return {input_rank->output(0)};
}
auto input_shape = make_shared<v3::ShapeOf>(input, ov::element::i32);
auto unsqueeze_input_rank = make_shared<v3::ShapeOf>(input_shape, ov::element::i32);
auto input_rank = make_shared<v0::Squeeze>(unsqueeze_input_rank);
Expand Down
42 changes: 42 additions & 0 deletions tests/layer_tests/tensorflow_tests/test_tf_Rank.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)

0 comments on commit 7d35d2a

Please sign in to comment.