Align handling of flaky arm unittests

backends/arm/test/ops/test_bmm.py

@@ -6,7 +6,7 @@
 
 import unittest
 
-from typing import Tuple
+from typing import Callable, Tuple
 
 import pytest
 
@@ -16,39 +16,37 @@
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from parameterized import parameterized
 
-torch.manual_seed(1)
-
 
 class TestBMM(unittest.TestCase):
     """Tests Batch MatMul"""
 
     class BMM(torch.nn.Module):
-        test_parameters = [
-            (torch.rand(2, 1, 1), torch.rand(2, 1, 1)),
-            (torch.rand(5, 3, 5), torch.rand(5, 5, 2)),
-            (torch.ones(1, 55, 3), torch.ones(1, 3, 44)),
-            (10000 * torch.randn(10, 1, 10), torch.randn(10, 10, 5)),
-            (-10 * torch.randn(2, 32, 64), 5 + 5 * torch.randn(2, 64, 32)),
+        test_data_generators = [
+            lambda: (torch.rand(2, 1, 1), torch.rand(2, 1, 1)),
+            lambda: (torch.rand(5, 3, 5), torch.rand(5, 5, 2)),
+            lambda: (torch.ones(1, 55, 3), torch.ones(1, 3, 44)),
+            lambda: (10000 * torch.randn(10, 1, 10), torch.randn(10, 10, 5)),
+            lambda: (-10 * torch.randn(2, 32, 64), 5 + 5 * torch.randn(2, 64, 32)),
         ]
 
         def forward(self, x, y):
             return torch.bmm(x, y)
 
     class MatMul(torch.nn.Module):
-        test_parameters = [
-            (torch.rand(2, 3, 5), torch.rand(2, 5, 2)),
-            (torch.rand(1, 2, 3, 5), torch.rand(1, 2, 5, 2)),
+        test_data_generators = [
+            lambda: (torch.rand(2, 3, 5), torch.rand(2, 5, 2)),
+            lambda: (torch.rand(1, 2, 3, 5), torch.rand(1, 2, 5, 2)),
         ]
 
         def forward(self, x, y):
             return torch.matmul(x, y)
 
     class BMMSingleInput(torch.nn.Module):
-        test_parameters = [
-            (torch.rand(20, 3, 3),),
-            (torch.rand(2, 128, 128),),
-            (10000 * torch.randn(4, 25, 25),),
-            (5 + 5 * torch.randn(3, 64, 64),),
+        test_data_generators = [
+            lambda: (torch.rand(20, 3, 3),),
+            lambda: (torch.rand(2, 128, 128),),
+            lambda: (10000 * torch.randn(4, 25, 25),),
+            lambda: (5 + 5 * torch.randn(3, 64, 64),),
         ]
 
         def forward(self, x):
@@ -120,67 +118,74 @@ def _test_bmm_ethosu_BI_pipeline(
         if conftest.is_option_enabled("corstone_fvp"):
             tester.run_method_and_compare_outputs(inputs=test_data, qtol=1)
 
-    @parameterized.expand(BMM.test_parameters)
-    def test_bmm_tosa_MI(self, operand1: torch.Tensor, operand2: torch.Tensor):
-        test_data = (operand1, operand2)
+    @parameterized.expand(BMM.test_data_generators)
+    def test_bmm_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_tosa_MI_pipeline(self.BMM(), test_data)
 
-    @parameterized.expand(BMMSingleInput.test_parameters)
-    def test_bmm_single_input_tosa_MI(self, operand1: torch.Tensor):
-        test_data = (operand1,)
+    @parameterized.expand(BMMSingleInput.test_data_generators)
+    def test_bmm_single_input_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_tosa_MI_pipeline(self.BMMSingleInput(), test_data)
 
-    @parameterized.expand(MatMul.test_parameters)
-    def test_matmul_tosa_MI(self, operand1: torch.Tensor, operand2: torch.Tensor):
-        test_data = (operand1, operand2)
+    @parameterized.expand(MatMul.test_data_generators)
+    def test_matmul_tosa_MI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_tosa_MI_pipeline(self.MatMul(), test_data)
 
-    @parameterized.expand(MatMul.test_parameters)
-    def test_matmul_tosa_BI(self, operand1: torch.Tensor, operand2: torch.Tensor):
-        test_data = (operand1, operand2)
+    @parameterized.expand(MatMul.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
+    def test_matmul_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_tosa_BI_pipeline(self.MatMul(), test_data)
 
-    @parameterized.expand(BMM.test_parameters)
-    def test_bmm_tosa_BI(self, operand1: torch.Tensor, operand2: torch.Tensor):
-        test_data = (operand1, operand2)
+    @parameterized.expand(BMM.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
+    def test_bmm_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_tosa_BI_pipeline(self.BMM(), test_data)
 
-    @parameterized.expand(BMMSingleInput.test_parameters)
-    def test_bmm_single_input_tosa_BI(self, operand1: torch.Tensor):
-        test_data = (operand1,)
+    @parameterized.expand(BMMSingleInput.test_data_generators)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
+    def test_bmm_single_input_tosa_BI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_tosa_BI_pipeline(self.BMMSingleInput(), test_data)
 
-    @parameterized.expand(BMM.test_parameters)
+    @parameterized.expand(BMM.test_data_generators)
     @pytest.mark.corstone_fvp
     @unittest.expectedFailure
-    def test_bmm_u55_BI_xfails(self, operand1: torch.Tensor, operand2: torch.Tensor):
-        test_data = (operand1, operand2)
+    def test_bmm_u55_BI_xfails(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_ethosu_BI_pipeline(
             self.BMM(), common.get_u55_compile_spec(), test_data
         )
 
-    @parameterized.expand(BMM.test_parameters)
+    @parameterized.expand(BMM.test_data_generators)
     @pytest.mark.corstone_fvp
-    def test_bmm_u85_BI(self, operand1: torch.Tensor, operand2: torch.Tensor):
-        test_data = (operand1, operand2)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
+    def test_bmm_u85_BI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_ethosu_BI_pipeline(
             self.BMM(), common.get_u85_compile_spec(), test_data
         )
 
     # Expected to fail with error: Warning, unsupported fusing of TOSA Rescale previous operator is of type: Memcpy
-    @parameterized.expand(BMMSingleInput.test_parameters)
+    @parameterized.expand(BMMSingleInput.test_data_generators)
     @pytest.mark.corstone_fvp
     @unittest.expectedFailure
-    def test_bmm_single_input_u55_BI_xfails(self, operand1: torch.Tensor):
-        test_data = (operand1,)
+    def test_bmm_single_input_u55_BI_xfails(
+        self, test_data_generator: Callable[[], Tuple]
+    ):
+        test_data = test_data_generator()
         self._test_bmm_ethosu_BI_pipeline(
             self.BMMSingleInput(), common.get_u55_compile_spec(), test_data
         )
 
-    @parameterized.expand(BMMSingleInput.test_parameters)
+    @parameterized.expand(BMMSingleInput.test_data_generators)
     @pytest.mark.corstone_fvp
-    def test_bmm_single_input_u85_BI(self, operand1: torch.Tensor):
-        test_data = (operand1,)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLETORCH-534)
+    def test_bmm_single_input_u85_BI(self, test_data_generator: Callable[[], Tuple]):
+        test_data = test_data_generator()
         self._test_bmm_ethosu_BI_pipeline(
             self.BMMSingleInput(), common.get_u85_compile_spec(), test_data
         )

backends/arm/test/ops/test_conv1d.py

@@ -6,7 +6,7 @@
 
 import unittest
 
-from typing import List, Optional, Tuple, Union
+from typing import List, Tuple, Union
 
 import pytest
 
@@ -25,7 +25,6 @@ class Conv1d(torch.nn.Module):
 
     def __init__(
         self,
-        inputs: Optional[torch.Tensor] = None,
         length=8,
         nbr_conv=1,  # Number of chained convs
         in_channels: Union[List, int, None] = None,
@@ -75,11 +74,10 @@ def __init__(
         if not isinstance(padding_mode, List):
             padding_mode = [padding_mode]
 
-        # Generate test data if not provided
-        if inputs is None:
-            self.inputs = (torch.randn(batches, in_channels[0], length).to(dtype),)
-        else:
-            self.inputs = (inputs,)
+        self.batches = batches
+        self.in_channels = in_channels
+        self.length = length
+        self.dtype = dtype
 
         # Build chain of convs
         for i in range(self.nbr_convs):
@@ -100,7 +98,9 @@ def __init__(
             )
 
     def get_inputs(self):
-        return self.inputs
+        return (
+            torch.randn(self.batches, self.in_channels[0], self.length).to(self.dtype),
+        )
 
     def forward(self, x):
         for i in range(self.nbr_convs):

backends/arm/test/ops/test_conv2d.py

@@ -6,7 +6,7 @@
 
 import unittest
 
-from typing import List, Optional, Tuple, Union
+from typing import List, Tuple, Union
 
 import pytest
 
@@ -25,7 +25,6 @@ class Conv2d(torch.nn.Module):
 
     def __init__(
         self,
-        inputs: Optional[torch.Tensor] = None,
         height=8,
         width=8,
         nbr_conv=1,  # Number of chained convs
@@ -76,13 +75,11 @@ def __init__(
         if not isinstance(padding_mode, List):
             padding_mode = [padding_mode]
 
-        # Generate test data if not provided
-        if inputs is None:
-            self.inputs = (
-                torch.randn(batches, in_channels[0], height, width).to(dtype),
-            )
-        else:
-            self.inputs = (inputs,)
+        self.batches = batches
+        self.in_channels = in_channels
+        self.height = height
+        self.width = width
+        self.dtype = dtype
 
         # Build chain of convs
         for i in range(self.nbr_convs):
@@ -103,7 +100,11 @@ def __init__(
             )
 
     def get_inputs(self):
-        return self.inputs
+        return (
+            torch.randn(self.batches, self.in_channels[0], self.height, self.width).to(
+                self.dtype
+            ),
+        )
 
     def forward(self, x):
         for i in range(self.nbr_convs):

backends/arm/test/ops/test_conv_combos.py

@@ -353,8 +353,7 @@ def test_block_bottleneck_residual_tosa_MI(self):
         model = ComboBlockBottleneckResidual()
         self._test_conv_combo_tosa_MI_pipeline(model, model.get_inputs())
 
-    # TODO: Investigate flakyness (MLTORCH-307)
-    @unittest.skip(reason="Skiped due to flakyness (MLTORCH-307)")
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLTORCH-307)
     def test_block_bottleneck_residual_tosa_BI(self):
         model = ComboBlockBottleneckResidual()
         self._test_conv_combo_tosa_BI_pipeline(model, model.get_inputs())

backends/arm/test/ops/test_depthwise_conv.py

@@ -252,8 +252,8 @@ def _test_dw_conv_ethos_BI_pipeline(
     def test_dw_conv_tosa_MI(self, test_name: str, model: torch.nn.Module):
         self._test_dw_conv_tosa_MI_pipeline(model, model.get_inputs())
 
-    # TODO: Investigate flakyness (MLTORCH-307)
     @parameterized.expand(testsuite_conv1d + testsuite_conv2d)
+    @pytest.mark.flaky  # TODO: Investigate flakyness (MLTORCH-307)
     def test_dw_conv_tosa_BI(self, test_name: str, model: torch.nn.Module):
         self._test_dw_conv_tosa_BI_pipeline(model, model.get_inputs())
 

backends/arm/test/ops/test_layer_norm.py

@@ -109,7 +109,7 @@ def _test_layernorm_tosa_BI_pipeline(
             .partition()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .to_executorch()
-            .run_method_and_compare_outputs(inputs=test_data)
+            .run_method_and_compare_outputs(qtol=1, inputs=test_data)
         )
 
     def _test_layernorm_ethosu_BI_pipeline(