[TF FE] Support complex tensors for OnesLike operation

src/frontends/tensorflow_common/src/op/ones_like.cpp

@@ -19,27 +19,56 @@ namespace tensorflow {
 namespace op {
 
 OutputVector translate_ones_like_op(const NodeContext& node) {
-    default_op_checks(node, 1, {"OnesLike"});
+    default_op_checks(node, 1, {"OnesLike"}, true);
     auto x = node.get_input(0);
-    Output<Node> shape_of = make_shared<v3::ShapeOf>(x, element::i32);
-    auto one_const = create_same_type_const_scalar<int32_t>(x, 1);
-
-    // in case of x to be scalar, we need handle it more specifically
-    // since Broadcast supports only broadcasting to rank greater 0
-    // we have to introduce extra dimension for input scalar case
-    auto one_dim = make_shared<v0::Constant>(element::i32, Shape{1}, 1);
-    shape_of = make_shared<v0::Concat>(OutputVector{one_dim, shape_of}, 0);
-
-    // create a tensor of zeros of shape with extra dimension
-    Output<Node> ones_like = make_shared<v3::Broadcast>(one_const, shape_of);
-    // remove extra dimension by squeezing
-    auto zero_dim_ind = make_shared<v0::Constant>(element::i32, Shape{1}, 0);
-    ones_like = make_shared<v0::Squeeze>(ones_like, zero_dim_ind);
-
-    set_node_name(node.get_name(), ones_like.get_node_shared_ptr());
-    return {ones_like};
-}
+    auto complex_type_mark_x = as_type_ptr<ComplexTypeMark>(x.get_node_shared_ptr());
+    if (complex_type_mark_x) {
+        x = complex_type_mark_x->input_value(0);
+        auto gather_index_real = make_shared<v0::Constant>(element::i32, Shape{}, 0);
+        auto minus_one = make_shared<v0::Constant>(element::i32, Shape{1}, -1);
+        auto x_real = make_shared<v8::Gather>(x, gather_index_real, minus_one)->output(0);
+
+        Output<Node> shape_of_real = make_shared<v3::ShapeOf>(x_real, element::i32);
+        auto one_const = create_same_type_const_scalar<int32_t>(x_real, 1);
+        auto one_dim = make_shared<v0::Constant>(element::i32, Shape{1}, 1);
+        shape_of_real = make_shared<v0::Concat>(OutputVector{one_dim, shape_of_real}, 0);
+        // create a tensor of zeros of shape with extra dimension
+        Output<Node> ones_like_one = make_shared<v3::Broadcast>(one_const, shape_of_real);
+        // remove extra dimension by squeezing
+        auto zero_dim_ind = make_shared<v0::Constant>(element::i32, Shape{1}, 0);
+        ones_like_one = make_shared<v0::Squeeze>(ones_like_one, zero_dim_ind);
+
+        // For zeroes:
+        auto zero_const = create_same_type_const_scalar<int32_t>(x_real, 0);
+        // create a tensor of zeros of shape with extra dimension
+        Output<Node> ones_like_zero = make_shared<v3::Broadcast>(zero_const, shape_of_real);
+        // remove extra dimension by squeezing
+        ones_like_zero = make_shared<v0::Squeeze>(ones_like_zero, zero_dim_ind);
 
+        auto ones_like_output = make_shared<v0::Concat>(OutputVector{ones_like_one, ones_like_zero}, -1);
+        auto ones_like_output_complex =
+            make_shared<ComplexTypeMark>(ones_like_output, ones_like_complex->get_complex_part_type())
+                set_node_name(node.get_name(), ones_like_output.get_node_shared_ptr());
+        return {ones_like_output_complex};
+    } else {
+        Output<Node> shape_of = make_shared<v3::ShapeOf>(x, element::i32);
+        auto one_const = create_same_type_const_scalar<int32_t>(x, 1);
+
+        // in case of x to be scalar, we need handle it more specifically
+        // since Broadcast supports only broadcasting to rank greater 0
+        // we have to introduce extra dimension for input scalar case
+        auto one_dim = make_shared<v0::Constant>(element::i32, Shape{1}, 1);
+        shape_of = make_shared<v0::Concat>(OutputVector{one_dim, shape_of}, 0);
+
+        // create a tensor of zeros of shape with extra dimension
+        Output<Node> ones_like = make_shared<v3::Broadcast>(one_const, shape_of);
+        // remove extra dimension by squeezing
+        auto zero_dim_ind = make_shared<v0::Constant>(element::i32, Shape{1}, 0);
+        ones_like = make_shared<v0::Squeeze>(ones_like, zero_dim_ind);
+        set_node_name(node.get_name(), ones_like.get_node_shared_ptr());
+        return {ones_like};
+    }
+}
 }  // namespace op
 }  // namespace tensorflow
 }  // namespace frontend

tests/layer_tests/tensorflow_tests/test_tf_OnesLike.py

@@ -43,3 +43,47 @@ def test_ones_like(self, params, ie_device, precision, ir_version, temp_dir,
         self._test(*self.create_ones_like_net(**params),
                    ie_device, precision, ir_version, temp_dir=temp_dir,
                    use_legacy_frontend=use_legacy_frontend)
+
+class TestComplexOnesLike(CommonTFLayerTest):
+    def _prepare_input(self, inputs_info):
+        rng = np.random.default_rng()
+        assert 'x_real:0' in inputs_info
+        assert 'x_imag:0' in inputs_info
+        x_real_shape = inputs_info['x_real:0']
+        x_imag_shape = inputs_info['x_imag:0']
+        inputs_data = {}
+        inputs_data['x_real:0'] = 4 * rng.random(x_real_shape).astype(self.x_type) - 2
+        inputs_data['x_imag:0'] = 4 * rng.random(x_imag_shape).astype(self.x_type) - 2
+        return inputs_data
+
+    def create_complex_ones_like_net(self, x_shape, x_type):
+        self.x_type = x_type
+        tf.compat.v1.reset_default_graph()
+        # Create the graph and model
+        with tf.compat.v1.Session() as sess:
+            x_real = tf.compat.v1.placeholder(tf.dtypes.as_dtype(x_type), x_shape, 'x_real')
+            x_imag = tf.compat.v1.placeholder(tf.dtypes.as_dtype(x_type), x_shape, 'x_imag')
+            x_complex = tf.raw_ops.Complex(real=x_real, imag=x_imag)
+            ones_like = tf.raw_ops.OnesLike(x=x_complex)
+            real = tf.raw_ops.Real(input=ones_like)
+            img = tf.raw_ops.Imag(input=ones_like)
+            tf.compat.v1.global_variables_initializer()
+            tf_net = sess.graph_def
+
+        return tf_net, None
+
+    test_data_basic = [
+        dict(x_shape=[], x_type=np.float32),
+        dict(x_shape=[2], x_type=np.int32),
+        dict(x_shape=[2, 3, 4], x_type=np.float32),
+        dict(x_shape=[1, 4, 3, 1], x_type=np.int32),
+    ]
+
+    @pytest.mark.parametrize("params", test_data_basic)
+    @pytest.mark.precommit_tf_fe
+    @pytest.mark.nightly
+    def test_complex_ones_like(self, params, ie_device, precision, ir_version, temp_dir,
+                       use_legacy_frontend):
+        self._test(*self.create_complex_ones_like_net(**params),
+                   ie_device, precision, ir_version, temp_dir=temp_dir,
+                   use_legacy_frontend=use_legacy_frontend)