[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 <roman.kazantsev@intel.com>
This commit is contained in:
Luca Tam 2024-03-12 22:49:24 +01:00 committed by GitHub
parent c1564dd90d
commit 81e236f51e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 63 additions and 1 deletions

View File

@ -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<ComplexTypeMark>(input.get_node_shared_ptr());
// retrive attribute of the output type
auto out_type = node.get_attribute<element::Type>("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<v3::ShapeOf>(input, out_type);
auto axis = make_shared<v0::Constant>(element::i32, Shape{}, 0);
auto complex_size = make_shared<v1::ReduceProd>(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<v1::Divide>(complex_size, make_shared<v0::Constant>(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<v0::Constant>(element::i32, Shape{1}, 0);
input = make_shared<v0::Unsqueeze>(input, const_zero);
@ -33,6 +52,7 @@ ov::OutputVector translate_size_op(const NodeContext& node) {
auto shape_of = make_shared<v3::ShapeOf>(input, out_type);
auto axis = make_shared<v0::Constant>(element::i32, Shape{}, 0);
auto size = make_shared<v1::ReduceProd>(shape_of, axis);
set_node_name(node.get_name(), size);
return {size};
}

View File

@ -43,3 +43,45 @@ class TestSize(CommonTFLayerTest):
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)