Fix AddBias Tests and NHCW logic

Pranshu-S · Mar 17, 2024 · f730b55 · f730b55
1 parent d7ed7d6
commit f730b55
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 20 deletions.
diff --git a/src/frontends/tensorflow_common/src/op/bias_add.cpp b/src/frontends/tensorflow_common/src/op/bias_add.cpp
@@ -47,11 +47,14 @@ OutputVector translate_bias_add_op(const NodeContext& node) {
         TENSORFLOW_OP_VALIDATION(node,
                                  value_shape.rank().is_static(),
                                  "Value of dynamic rank for BiasAdd in NCHW layout is not supported.");
-        auto value_rank = complex_type_inputs ? value_shape.rank().get_length() - 1 : value_shape.rank().get_length();
+        auto value_rank = complex_type_inputs ? value_shape.rank().get_length() : value_shape.rank().get_length();
 
         std::vector<int64_t> axes_unsqueeze;
         for (int64_t dim_ind = 0; dim_ind < value_rank; ++dim_ind) {
-            if (dim_ind != 1) {
+            if (!complex_type_inputs && dim_ind != 1) {
+                axes_unsqueeze.push_back(dim_ind);
+            }
+            if (complex_type_inputs && dim_ind != 2){
                 axes_unsqueeze.push_back(dim_ind);
             }
         }

diff --git a/tests/layer_tests/tensorflow_tests/test_tf_BiasAdd.py b/tests/layer_tests/tensorflow_tests/test_tf_BiasAdd.py
@@ -161,27 +161,29 @@ 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']
+        assert 'y_real:0' in inputs_info
+        assert 'y_imag:0' in inputs_info
+        x_shape = inputs_info['x_real:0']
+        y_shape = inputs_info['y_real:0']
         inputs_data = {}
-        inputs_data['x_real:0'] = 4 * rng.random(x_real_shape).astype(np.float64) - 2
-        inputs_data['x_imag:0'] = 4 * rng.random(x_imag_shape).astype(np.float64) - 2
+
+        inputs_data['x_real:0'] = 4 * rng.random(x_shape).astype(np.float64) - 2
+        inputs_data['x_imag:0'] = 4 * rng.random(x_shape).astype(np.float64) - 2
+
+        inputs_data['y_real:0'] = 4 * rng.random(y_shape).astype(np.float64) - 2
+        inputs_data['y_imag:0'] = 4 * rng.random(y_shape).astype(np.float64) - 2
+
         return inputs_data
 
-    def create_complex_bias_add_net(self, shape, data_format, ir_version, use_legacy_frontend, output_type=tf.float64):
+    def create_complex_bias_add_net(self, input_shape, bias_shape, data_format, ir_version, use_legacy_frontend, output_type=tf.float64):
         tf.compat.v1.reset_default_graph()
 
         with tf.compat.v1.Session() as sess:
-            x_real_shape = shape.copy()
-            x_imag_shape = shape.copy()
-
-            x_real = tf.compat.v1.placeholder(output_type, x_real_shape, 'x_real')
-            x_imag = tf.compat.v1.placeholder(output_type, x_imag_shape, 'x_imag')
+            x_real = tf.compat.v1.placeholder(output_type, input_shape, 'x_real')
+            x_imag = tf.compat.v1.placeholder(output_type, input_shape, 'x_imag')
 
-            constant_value_real = np.random.randint(-256, 256, x_real_shape[-1]).astype(output_type.as_numpy_dtype())
-            constant_value_imag = np.random.randint(-256, 256, x_imag_shape[-1]).astype(output_type.as_numpy_dtype())
-            y_real = tf.constant(constant_value_real)
-            y_imag = tf.constant(constant_value_imag)
+            y_real = tf.compat.v1.placeholder(output_type, bias_shape, 'y_real')
+            y_imag = tf.compat.v1.placeholder(output_type, bias_shape, 'y_imag')
 
             complex_input = tf.complex(x_real, x_imag)
             complex_bias = tf.complex(y_real, y_imag)
@@ -195,10 +197,10 @@ def create_complex_bias_add_net(self, shape, data_format, ir_version, use_legacy
         return tf_net, None
 
     test_data_2D = [
-        dict(shape=[1, 1], data_format="NHWC"),
-        dict(shape=[1, 224], data_format="NHWC"),
-        dict(shape=[1, 1], data_format="NCHW"),
-        dict(shape=[1, 224], data_format="NCHW"),
+        dict(shape=[1, 1], bias_shape=[1], data_format="NHWC"),
+        dict(shape=[3, 2, 7], bias_shape=[7], data_format="NHWC"),
+        dict(shape=[3, 2, 7, 10], bias_shape=[2], data_format="NCHW"),
+        dict(shape=[7, 6, 4, 5], bias_shape=[6], data_format="NCHW"),
     ]
 
     @pytest.mark.parametrize("params", test_data_2D)