Arm backend: Mark test in test_bmm.py as flaky (pytorch#14748)

Add a new argument to `common.parametrize`, `flakies`, which selects
which parametrized test cases to mark as flaky. With this new argument,
mark the test test_bmm.py::test_bmm_vgf_FP_single_input[rand_big_1] as
flaky.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218

Signed-off-by: Martin Lindström <[email protected]>
Co-authored-by: Martin Lindström <[email protected]>

Author: martinlsm

backends/arm/test/common.py

@@ -227,6 +227,7 @@ def parametrize(
     test_data: dict[str, Any],
     xfails: dict[str, xfail_type] | None = None,
     strict: bool = True,
+    flakies: dict[str, int] | None = None,
 ):
     """
     Custom version of pytest.mark.parametrize with some syntatic sugar and added xfail functionality
@@ -237,12 +238,17 @@ def parametrize(
     """
     if xfails is None:
         xfails = {}
+    if flakies is None:
+        flakies = {}
 
     def decorator_func(func):
         """Test data is transformed from a dict of (id, data) pairs to a list of pytest params to work with the native pytests parametrize function"""
         pytest_testsuite = []
         for id, test_parameters in test_data.items():
-            if id in xfails:
+            if id in flakies:
+                # Mark this parameter as flaky with given reruns
+                marker = (pytest.mark.flaky(reruns=flakies[id]),)
+            elif id in xfails:
                 xfail_info = xfails[id]
                 reason = ""
                 raises = None

backends/arm/test/ops/test_bmm.py

@@ -146,7 +146,11 @@ def test_bmm_vgf_FP(test_data: input_t1):
     pipeline.run()
 
 
-@common.parametrize("test_data", BMMSingleInput.test_data_generators)
+@common.parametrize(
+    "test_data",
+    BMMSingleInput.test_data_generators,
+    flakies={"rand_big_1": 3},
+)
 @common.SkipIfNoModelConverter
 def test_bmm_vgf_FP_single_input(test_data: input_t1):
     pipeline = VgfPipeline[input_t1](