Skip to content

Commit 6929f7e

Browse files
11happyrkazants
andauthored
[TF FE] Feat: implement complex type support for select (#28677)
**Overview:** - This pull request fixes #23243. **Testing:** - Tested the implementation & verified the results. ![Screenshot from 2025-01-25 14-21-47](https://github.com/user-attachments/assets/0cab7b27-727e-4085-8719-fda318eade16) - No dependencies on other pull requests. - TODO: Improve the test cases. **CC:** - @rkazants --------- Signed-off-by: 11happy <soni5happy@gmail.com> Co-authored-by: Roman Kazantsev <roman.kazantsev@intel.com>
1 parent 33a616d commit 6929f7e

File tree

2 files changed

+76
-2
lines changed

2 files changed

+76
-2
lines changed

src/frontends/tensorflow_common/src/op/select.cpp

+25-2
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
#include "openvino/op/select.hpp"
66

77
#include "common_op_table.hpp"
8+
#include "helper_ops/complex_type_mark.hpp"
89
#include "openvino/op/broadcast.hpp"
910
#include "openvino/op/concat.hpp"
1011
#include "openvino/op/constant.hpp"
@@ -54,10 +55,25 @@ OutputVector translate_select_op(const NodeContext& node) {
5455
// 1. Either the same shape (in which case the select is elementwise), or
5556
// 2. condition must be Rank 1 and match over the first dimension, or
5657
// 3. condition is scalar
57-
default_op_checks(node, 3, {"Select", "SELECT"});
58+
default_op_checks(node, 3, {"Select", "SELECT"}, true);
5859
auto condition = node.get_input(0);
5960
auto x = node.get_input(1);
6061
auto y = node.get_input(2);
62+
auto complex_type_mark_x = as_type_ptr<ComplexTypeMark>(x.get_node_shared_ptr());
63+
auto complex_type_mark_y = as_type_ptr<ComplexTypeMark>(y.get_node_shared_ptr());
64+
65+
auto is_complex = (complex_type_mark_x || complex_type_mark_y);
66+
element::Type complex_part_type;
67+
68+
if (complex_type_mark_x) {
69+
x = complex_type_mark_x->input_value(0);
70+
complex_part_type = complex_type_mark_x->get_complex_part_type();
71+
}
72+
73+
if (complex_type_mark_y) {
74+
y = complex_type_mark_y->input_value(0);
75+
complex_part_type = complex_type_mark_y->get_complex_part_type();
76+
}
6177

6278
// compute number of dimensions to unsqueeze the condition
6379
auto cond_rank = compute_subgraph_scalar_rank(condition, element::i32);
@@ -78,7 +94,14 @@ OutputVector translate_select_op(const NodeContext& node) {
7894
auto const_0 = make_shared<v0::Constant>(element::i32, Shape{1}, 0);
7995
prep_cond = make_shared<v0::Squeeze>(prep_cond, const_0);
8096

81-
return translate_select_base_op(node, prep_cond, x, y);
97+
auto result = translate_select_base_op(node, prep_cond, x, y);
98+
if (is_complex) {
99+
auto complex_result = make_shared<ComplexTypeMark>(result[0].get_node_shared_ptr(), complex_part_type);
100+
return {complex_result->output(0)};
101+
102+
} else {
103+
return result;
104+
}
82105
}
83106
} // namespace op
84107
} // namespace tensorflow

tests/layer_tests/tensorflow_tests/test_tf_Select.py

+51
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,54 @@ def test_select_basic(self, params, ie_device, precision, ir_version, temp_dir,
5252
self._test(*self.create_select_net(**params),
5353
ie_device, precision, ir_version, temp_dir=temp_dir,
5454
use_legacy_frontend=use_legacy_frontend)
55+
56+
class TestComplexSelect(CommonTFLayerTest):
57+
def _prepare_input(self, inputs_info):
58+
rng = np.random.default_rng()
59+
assert 'cond:0' in inputs_info, "Test error: inputs_info must contain `cond`"
60+
assert 'x_real:0' in inputs_info, "Test error: inputs_info must contain `x_real`"
61+
assert 'x_imag:0' in inputs_info, "Test error: inputs_info must contain `x_imag`"
62+
assert 'y_real:0' in inputs_info, "Test error: inputs_info must contain `y_real`"
63+
assert 'y_imag:0' in inputs_info, "Test error: inputs_info must contain `y_imag`"
64+
cond_shape = inputs_info['cond:0']
65+
inputs_data = {}
66+
inputs_data['cond:0'] = np.random.randint(0, 2, cond_shape).astype(bool)
67+
for part in ['x_real:0', 'x_imag:0', 'y_real:0', 'y_imag:0']:
68+
inputs_data[part] = 4 * rng.random(inputs_info[part]).astype(np.float32) - 2
69+
return inputs_data
70+
71+
def create_complex_select_net(self, cond_shape, x_shape, y_shape):
72+
tf.compat.v1.reset_default_graph()
73+
# Create the graph and model
74+
with tf.compat.v1.Session() as sess:
75+
cond = tf.compat.v1.placeholder(tf.bool, cond_shape, 'cond')
76+
x_real = tf.compat.v1.placeholder(tf.float32, x_shape, 'x_real')
77+
x_imag = tf.compat.v1.placeholder(tf.float32, x_shape, 'x_imag')
78+
y_real = tf.compat.v1.placeholder(tf.float32, y_shape, 'y_real')
79+
y_imag = tf.compat.v1.placeholder(tf.float32, y_shape, 'y_imag')
80+
complex_x = tf.raw_ops.Complex(real=x_real, imag=x_imag)
81+
complex_y = tf.raw_ops.Complex(real=y_real, imag=y_imag)
82+
complex_select = tf.raw_ops.Select(condition=cond, x=complex_x, y=complex_y)
83+
tf.raw_ops.Real(input=complex_select)
84+
tf.raw_ops.Imag(input=complex_select)
85+
tf.compat.v1.global_variables_initializer()
86+
tf_net = sess.graph_def
87+
return tf_net, None
88+
89+
test_data_basic = [
90+
dict(cond_shape=[], x_shape=[], y_shape=[]),
91+
dict(cond_shape=[], x_shape=[3, 2, 4], y_shape=[3, 2, 4]),
92+
dict(cond_shape=[2, 3, 4], x_shape=[2, 3, 4], y_shape=[2, 3, 4]),
93+
]
94+
95+
@pytest.mark.parametrize("params", test_data_basic)
96+
@pytest.mark.precommit
97+
@pytest.mark.nightly
98+
99+
def test_complex_select(self, params, ie_device, precision, ir_version, temp_dir,
100+
use_legacy_frontend):
101+
if use_legacy_frontend:
102+
pytest.skip("Select tests are not passing for the legacy frontend.")
103+
self._test(*self.create_complex_select_net(**params),
104+
ie_device, precision, ir_version, temp_dir=temp_dir,
105+
use_legacy_frontend=use_legacy_frontend)

0 commit comments

Comments
 (0)