Fix failing test

Author: Pranshu-S

tests/layer_tests/tensorflow2_keras_tests/test_tf2_map_fn.py

@@ -6,38 +6,44 @@
 
 from common.tf2_layer_test_class import CommonTF2LayerTest
 
-
-class MapFNLayer(tf.keras.layers.Layer):
-    def __init__(self, fn, input_type, fn_output_signature, back_prop):
-        super(MapFNLayer, self).__init__()
-        self.fn = fn
-        self.input_type = input_type
-        self.fn_output_signature = fn_output_signature
-        self.back_prop = back_prop
-
-    def call(self, x):
-        return tf.map_fn(self.fn, x, dtype=self.input_type,
-                         fn_output_signature=self.fn_output_signature,
-                         back_prop=self.back_prop)
-
 def fn_1(x):
     return (x[0] * x[1] + x[2])
 
 def fn_2(x):
     return (x[0] + x[1] + x[2], x[0] - x[2] + x[1], 2 + x[2])
 
 def fn_3(x):
-    return (x[0] * x[1], x[0] + x[1])
+    return (x[0] * x[1])
 
 def fn_4(x):
-    return (x[0] * x[1] + x[2], x[0] + x[2] * x[1], 2 * x[2])
+    return (x[0] * x[1] + 2 * x[2])
 
 def fn_5(x):
-    return (x[0] * x[1] + x[2])
+    return (x[0] * x[1], x[0] + x[1])
 
 def fn_6(x):
+    return (x[0] * x[1] + x[2], x[0] + x[2] * x[1], 2 * x[2])
+
+def fn_7(x):
+    return (x[0] * x[1] + x[2])
+
+def fn_8(x):
     return (x[0] + x[1] + x[2], x[0] - x[2] + x[1], 2 + x[2])
 
+list_fns = [fn_1, fn_2, fn_3, fn_4, fn_5, fn_6, fn_7, fn_8]
+
+class MapFNLayer(tf.keras.layers.Layer):
+    def __init__(self, fn, input_type, fn_output_signature, back_prop):
+        super(MapFNLayer, self).__init__()
+        self.fn = list_fns[fn-1]
+        self.input_type = input_type
+        self.fn_output_signature = fn_output_signature
+        self.back_prop = back_prop
+
+    def call(self, x):
+        return tf.map_fn(self.fn, x, dtype=self.input_type,
+                         fn_output_signature=self.fn_output_signature,
+                         back_prop=self.back_prop)
 
 class TestMapFN(CommonTF2LayerTest):
     def create_map_fn_net(self, fn, input_type, fn_output_signature, back_prop,
@@ -57,10 +63,10 @@ def create_map_fn_net(self, fn, input_type, fn_output_signature, back_prop,
         return tf2_net, ref_net
 
     test_basic = [
-        dict(fn=fn_1, input_type=tf.float32,
+        dict(fn=1, input_type=tf.float32,
              fn_output_signature=tf.float32, back_prop=False,
              input_names=["x1", "x2", "x3"], input_shapes=[[2, 3, 4], [2, 3, 4], [2, 3, 4]]),
-        pytest.param(dict(fn=fn_2,
+        pytest.param(dict(fn=2,
                           input_type=tf.float32,
                           fn_output_signature=(tf.float32, tf.float32, tf.float32), back_prop=True,
                           input_names=["x1", "x2", "x3"],
@@ -77,10 +83,10 @@ def test_basic(self, params, ie_device, precision, ir_version, temp_dir, use_leg
                    **params)
 
     test_multiple_inputs = [
-        dict(fn=lambda x: x[0] * x[1], input_type=tf.float32,
+        dict(fn=3, input_type=tf.float32,
              fn_output_signature=tf.float32, back_prop=True,
              input_names=["x1", "x2"], input_shapes=[[2, 4], [2, 4]]),
-        dict(fn=lambda x: x[0] * x[1] + 2 * x[2], input_type=tf.float32,
+        dict(fn=4, input_type=tf.float32,
              fn_output_signature=tf.float32, back_prop=False,
              input_names=["x1", "x2", "x3"], input_shapes=[[2, 1, 3, 4],
                                                            [2, 1, 3, 4],
@@ -95,11 +101,11 @@ def test_multiple_inputs(self, params, ie_device, precision, ir_version, temp_di
                    **params)
 
     test_multiple_outputs = [
-        pytest.param(dict(fn=fn_3, input_type=tf.float32,
+        pytest.param(dict(fn=5, input_type=tf.float32,
                           fn_output_signature=(tf.float32, tf.float32), back_prop=True,
                           input_names=["x1", "x2"], input_shapes=[[2, 4], [2, 4]]),
                      marks=pytest.mark.xfail(reason="61587")),
-        pytest.param(dict(fn=fn_4,
+        pytest.param(dict(fn=6,
                           input_type=tf.float32,
                           fn_output_signature=(tf.float32, tf.float32, tf.float32), back_prop=True,
                           input_names=["x1", "x2", "x3"],
@@ -115,12 +121,12 @@ def test_multiple_outputs(self, params, ie_device, precision, ir_version, temp_d
                    **params)
 
     test_multiple_inputs_outputs_int32 = [
-        dict(fn=fn_5,
+        dict(fn=7,
              input_type=tf.int32,
              fn_output_signature=tf.int32, back_prop=True,
              input_names=["x1", "x2", "x3"],
              input_shapes=[[2, 1, 3], [2, 1, 3], [2, 1, 3]]),
-        pytest.param(dict(fn=fn_6,
+        pytest.param(dict(fn=8,
                           input_type=tf.int32,
                           fn_output_signature=(tf.int32, tf.int32, tf.int32), back_prop=True,
                           input_names=["x1", "x2", "x3"],