Avoid unnecessary slow compilation of Aesara functions in test_step

Author: ricardoV94

pymc/tests/test_step.py

@@ -60,6 +60,7 @@
 )
 from pymc.step_methods.mlda import extract_Q_estimate
 from pymc.tests.checks import close_to
+from pymc.tests.helpers import fast_unstable_sampling_mode
 from pymc.tests.models import (
     mv_simple,
     mv_simple_coarse,
@@ -175,20 +176,21 @@ def test_step_categorical(self, proposal):
 
 class TestMetropolisProposal:
     def test_proposal_choice(self):
-        _, model, _ = mv_simple()
-        with model:
-            initial_point = model.initial_point()
-            initial_point_size = sum(initial_point[n.name].size for n in model.value_vars)
+        with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+            _, model, _ = mv_simple()
+            with model:
+                initial_point = model.initial_point()
+                initial_point_size = sum(initial_point[n.name].size for n in model.value_vars)
 
-            s = np.ones(initial_point_size)
-            sampler = Metropolis(S=s)
-            assert isinstance(sampler.proposal_dist, NormalProposal)
-            s = np.diag(s)
-            sampler = Metropolis(S=s)
-            assert isinstance(sampler.proposal_dist, MultivariateNormalProposal)
-            s[0, 0] = -s[0, 0]
-            with pytest.raises(np.linalg.LinAlgError):
+                s = np.ones(initial_point_size)
                 sampler = Metropolis(S=s)
+                assert isinstance(sampler.proposal_dist, NormalProposal)
+                s = np.diag(s)
+                sampler = Metropolis(S=s)
+                assert isinstance(sampler.proposal_dist, MultivariateNormalProposal)
+                s[0, 0] = -s[0, 0]
+                with pytest.raises(np.linalg.LinAlgError):
+                    sampler = Metropolis(S=s)
 
     def test_mv_proposal(self):
         np.random.seed(42)
@@ -202,59 +204,60 @@ def test_mv_proposal(self):
 class TestCompoundStep:
     samplers = (Metropolis, Slice, HamiltonianMC, NUTS, DEMetropolis)
 
-    @pytest.mark.skipif(
-        aesara.config.floatX == "float32", reason="Test fails on 32 bit due to linalg issues"
-    )
     def test_non_blocked(self):
         """Test that samplers correctly create non-blocked compound steps."""
-        _, model = simple_2model_continuous()
-        with model:
-            for sampler in self.samplers:
-                assert isinstance(sampler(blocked=False), CompoundStep)
+        with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+            _, model = simple_2model_continuous()
+            with model:
+                for sampler in self.samplers:
+                    assert isinstance(sampler(blocked=False), CompoundStep)
 
-    @pytest.mark.skipif(
-        aesara.config.floatX == "float32", reason="Test fails on 32 bit due to linalg issues"
-    )
     def test_blocked(self):
-        _, model = simple_2model_continuous()
-        with model:
-            for sampler in self.samplers:
-                sampler_instance = sampler(blocked=True)
-                assert not isinstance(sampler_instance, CompoundStep)
-                assert isinstance(sampler_instance, sampler)
+        with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+            _, model = simple_2model_continuous()
+            with model:
+                for sampler in self.samplers:
+                    sampler_instance = sampler(blocked=True)
+                    assert not isinstance(sampler_instance, CompoundStep)
+                    assert isinstance(sampler_instance, sampler)
 
 
 class TestAssignStepMethods:
     def test_bernoulli(self):
         """Test bernoulli distribution is assigned binary gibbs metropolis method"""
         with Model() as model:
             Bernoulli("x", 0.5)
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, BinaryGibbsMetropolis)
 
     def test_normal(self):
         """Test normal distribution is assigned NUTS method"""
         with Model() as model:
             Normal("x", 0, 1)
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, NUTS)
 
     def test_categorical(self):
         """Test categorical distribution is assigned categorical gibbs metropolis method"""
         with Model() as model:
             Categorical("x", np.array([0.25, 0.75]))
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, BinaryGibbsMetropolis)
         with Model() as model:
             Categorical("y", np.array([0.25, 0.70, 0.05]))
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, CategoricalGibbsMetropolis)
 
     def test_binomial(self):
         """Test binomial distribution is assigned metropolis method."""
         with Model() as model:
             Binomial("x", 10, 0.5)
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, Metropolis)
 
     def test_normal_nograd_op(self):
@@ -274,7 +277,8 @@ def kill_grad(x):
             data = np.random.normal(size=(100,))
             Normal("y", mu=kill_grad(x), sigma=1, observed=data.astype(aesara.config.floatX))
 
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, Slice)
 
     def test_modify_step_methods(self):
@@ -286,15 +290,17 @@ def test_modify_step_methods(self):
 
         with Model() as model:
             Normal("x", 0, 1)
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert not isinstance(steps, NUTS)
 
         # add back nuts
         pm.STEP_METHODS = step_methods + [NUTS]
 
         with Model() as model:
             Normal("x", 0, 1)
-            steps = assign_step_methods(model, [])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                steps = assign_step_methods(model, [])
         assert isinstance(steps, NUTS)
 
 
@@ -1326,7 +1332,8 @@ def test_continuous_steps(self, step, step_kwargs):
             c1 = HalfNormal("c1")
             c2 = HalfNormal("c2")
 
-            assert [m.rvs_to_values[c1]] == step([c1], **step_kwargs).vars
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                assert [m.rvs_to_values[c1]] == step([c1], **step_kwargs).vars
             assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(
                 step([c1, c2], **step_kwargs).vars
             )
@@ -1343,7 +1350,8 @@ def test_discrete_steps(self, step, step_kwargs):
             d1 = Bernoulli("d1", p=0.5)
             d2 = Bernoulli("d2", p=0.5)
 
-            assert [m.rvs_to_values[d1]] == step([d1], **step_kwargs).vars
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                assert [m.rvs_to_values[d1]] == step([d1], **step_kwargs).vars
             assert {m.rvs_to_values[d1], m.rvs_to_values[d2]} == set(
                 step([d1, d2], **step_kwargs).vars
             )
@@ -1353,7 +1361,8 @@ def test_compound_step(self):
             c1 = HalfNormal("c1")
             c2 = HalfNormal("c2")
 
-            step1 = NUTS([c1])
-            step2 = NUTS([c2])
-            step = CompoundStep([step1, step2])
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                step1 = NUTS([c1])
+                step2 = NUTS([c2])
+                step = CompoundStep([step1, step2])
             assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(step.vars)