Test variants of the radon model function

Author: ricardoV94

tests/compile/function/test_types.py

@@ -1490,26 +1490,11 @@ def zerosumnormal(name, *, sigma=1.0, size, model_logp):
     return joined_inputs, [model_logp, model_dlogp]
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def radon_model():
     return create_radon_model()
 
 
-@pytest.fixture(scope="session")
-def radon_model_variants():
-    # Convert to list comp
-    return [
-        create_radon_model(
-            intercept_dist=intercept_dist,
-            sigma_dist=sigma_dist,
-            centered=centered,
-        )
-        for centered in (True, False)
-        for intercept_dist in ("normal", "lognormal")
-        for sigma_dist in ("halfnormal", "lognormal")
-    ]
-
-
 @pytest.mark.parametrize("mode", ["C", "C_VM", "NUMBA"])
 def test_radon_model_repeated_compile_benchmark(mode, radon_model, benchmark):
     joined_inputs, [model_logp, model_dlogp] = radon_model
@@ -1526,15 +1511,13 @@ def compile_and_call_once():
 
 
 @pytest.mark.parametrize("mode", ["C", "C_VM", "NUMBA"])
-def test_radon_model_variants_compile_benchmark(
-    mode, radon_model, radon_model_variants, benchmark
-):
+def test_radon_model_small_variants_compile_benchmark(mode, radon_model, benchmark):
     """Test compilation speed when a slightly variant of a function is compiled each time.
 
     This test more realistically simulates a use case where a model is recompiled
     multiple times with small changes, such as in an interactive environment.
 
-    NOTE: For this test to be meaningful on subsequent runs, the cache must be cleared
+    TODO: Change cache directory to force recompile. This must be done before pytensor is imported
     """
     joined_inputs, [model_logp, model_dlogp] = radon_model
     rng = np.random.default_rng(1)
@@ -1547,11 +1530,21 @@ def test_radon_model_variants_compile_benchmark(
     fn(x)
 
     def compile_and_call_once():
-        for joined_inputs, [model_logp, model_dlogp] in radon_model_variants:
-            fn = function(
-                [joined_inputs], [model_logp, model_dlogp], mode=mode, trust_input=True
-            )
-            fn(x)
+        for centered in (True, False):
+            for intercept_dist in ("normal", "lognormal"):
+                for sigma_dist in ("halfnormal", "lognormal"):
+                    joined_inputs, [model_logp, model_dlogp] = create_radon_model(
+                        intercept_dist=intercept_dist,
+                        sigma_dist=sigma_dist,
+                        centered=centered,
+                    )
+                    fn = function(
+                        [joined_inputs],
+                        [model_logp, model_dlogp],
+                        mode=mode,
+                        trust_input=True,
+                    )
+                    fn(x)
 
     benchmark.pedantic(compile_and_call_once, rounds=1, iterations=1)