Repository: pymc-devs/pymc

Issue:
Slice sampler does not converge in simple MvNormal prior model
This example is taken from the step_tests here: https://github.com/pymc-devs/pymc/blob/4d2f3a8d5f99fc773bc0a1aad6d51e4f3f047698/pymc/tests/test_step.py#L89

**Please provide a minimal, self-contained, and reproducible example.**
```python
import pymc as pm
import numpy as np

mu = np.array([-0.1, 0.5, 1.1])
p = np.array([[2.0, 0, 0], [0.05, 0.1, 0], [1.0, -0.05, 5.5]])
tau = np.dot(p, p.T)

with pm.Model() as model:
    pm.MvNormal( "x",  mu, tau=tau,)

expected_mean = mu
expected_std = np.diag(np.linalg.inv(tau)) ** 0.5

with model:
    trace = pm.sample(
        tune=10_000,
        draws=10_000,
        chains=1,
        step=pm.Slice(),
        return_inferencedata=False,
    )

# Very bad, the second mean seems to be consistently biased
print(expected_mean, trace["x"].mean(axis=0))
# [-0.1  0.5  1.1] [-0.14368236  1.76355742  1.11952911]

# Quite good
print(expected_std, trace["x"].std(axis=0)) 
# [ 0.56672996 10.00041321  0.18181818] [0.56335647 9.88682831 0.18468393]
```

This was also the case in V3: https://colab.research.google.com/drive/1Pw7a-1J-bnb4dPsPvxwuWH6TdFENYb5m?usp=sharing

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5816?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5816](https://codecov.io/gh/pymc-devs/pymc/pull/5816?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (29a3f48) into [main](https://codecov.io/gh/pymc-devs/pymc/commit/4d2f3a8d5f99fc773bc0a1aad6d51e4f3f047698?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (4d2f3a8) will **decrease** coverage by `25.07%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head 29a3f48 differs from pull request most recent head 545d4ff. Consider uploading reports for the commit 545d4ff to get more accurate results

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5816/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://codecov.io/gh/pymc-devs/pymc/pull/5816?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #5816       +/-   ##
===========================================
- Coverage   89.39%   64.32%   -25.08%     
===========================================
  Files          74       74               
  Lines       13767    13770        +3     
===========================================
- Hits        12307     8857     -3450     
- Misses       1460     4913     +3453     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5816?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/step\_methods/compound.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvY29tcG91bmQucHk=) | `86.66% <100.00%> (-6.52%)` | :arrow_down: |
| [pymc/step\_methods/slicer.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvc2xpY2VyLnB5) | `95.89% <100.00%> (+0.11%)` | :arrow_up: |
| [pymc/ode/utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9vZGUvdXRpbHMucHk=) | `15.62% <0.00%> (-84.38%)` | :arrow_down: |
| [pymc/smc/smc.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zbWMvc21jLnB5) | `18.47% <0.00%> (-77.92%)` | :arrow_down: |
| [pymc/gp/gp.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9ncC9ncC5weQ==) | `18.68% <0.00%> (-74.50%)` | :arrow_down: |
| [pymc/tuning/starting.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy90dW5pbmcvc3RhcnRpbmcucHk=) | `19.83% <0.00%> (-72.73%)` | :arrow_down: |
| [pymc/gp/cov.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9ncC9jb3YucHk=) | `29.15% <0.00%> (-68.94%)` | :arrow_down: |
| [pymc/distributions/mixture.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL21peHR1cmUucHk=) | `29.18% <0.00%> (-67.57%)` | :arrow_down: |
| [pymc/smc/sample\_smc.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zbWMvc2FtcGxlX3NtYy5weQ==) | `16.77% <0.00%> (-66.45%)` | :arrow_down: |
| [pymc/func\_utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9mdW5jX3V0aWxzLnB5) | `22.50% <0.00%> (-65.00%)` | :arrow_down: |
| ... and [37 more](https://codecov.io/gh/pymc-devs/pymc/pull/5816/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |

ricardoV94: It seems like this is how the sampler used to behave until this commit https://github.com/pymc-devs/pymc/commit/b988ba99a084a18b2c907ef2ee4aeb1f4ab56160

changed `right_bound = left_bound + w` into `right_bound = initial_point + w`
michaelosthege: > It seems like this is how the sampler used to behave until commit [b988ba9](https://github.com/pymc-devs/pymc/commit/b988ba99a084a18b2c907ef2ee4aeb1f4ab56160) changed `right_bound = left_bound + w` -> `right_bound = initial_point + w`

Wow that means the Slice sampler was broken since version 3.2.


ricardoV94: > Wow that means the Slice sampler was broken since version 3.2.


Should we backport the fix to V3?


michaelosthege: > Should we backport the fix to V3?

I'm not sure if it's worth it, but an exception telling users to upgrade to v4 should be easy enough. 
ricardoV94: It's the same amount of work to add an exception or change the 2 lines
ricardoV94: I should mention something in the Release Notes

Patch:
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 045e2ccff..8c896ae66 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -148,8 +148,7 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
 - Added the low level `compile_forward_sampling_function` method to compile the aesara function responsible for generating forward samples (see [#5759](https://github.com/pymc-devs/pymc/pull/5759)).
 - ...
 
-
-## Documentation
+### Documentation
 - Switched to the [pydata-sphinx-theme](https://pydata-sphinx-theme.readthedocs.io/en/latest/)
 - Updated our documentation tooling to use [MyST](https://myst-parser.readthedocs.io/en/latest/), [MyST-NB](https://myst-nb.readthedocs.io/en/latest/), sphinx-design, notfound.extension,
   sphinx-copybutton and sphinx-remove-toctrees.
@@ -157,7 +156,8 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
 - Restructured the documentation to facilitate learning paths
 - Updated API docs to document objects at the path users should use to import them
 
-### Internal changes
+### Maintenance
+- ⚠ Fixed old-time bug in Slice sampler that resulted in biased samples (see [#5816](https://github.com/pymc-devs/pymc/pull/5816)).
 - ⚠ PyMC now requires Scipy version `>= 1.4.1` (see [4857](https://github.com/pymc-devs/pymc/pull/4857)).
 - Removed float128 dtype support (see [#4514](https://github.com/pymc-devs/pymc/pull/4514)).
 - Logp method of `Uniform` and `DiscreteUniform` no longer depends on `pymc.distributions.dist_math.bound` for proper evaluation (see [#4541](https://github.com/pymc-devs/pymc/pull/4541)).
diff --git a/pymc/step_methods/compound.py b/pymc/step_methods/compound.py
index 9a0326096..d824045c7 100644
--- a/pymc/step_methods/compound.py
+++ b/pymc/step_methods/compound.py
@@ -33,6 +33,9 @@ class CompoundStep:
         for method in self.methods:
             if method.generates_stats:
                 self.stats_dtypes.extend(method.stats_dtypes)
+        self.name = (
+            f"Compound[{', '.join(getattr(m, 'name', 'UNNAMED_STEP') for m in self.methods)}]"
+        )
 
     def step(self, point):
         if self.generates_stats:
diff --git a/pymc/step_methods/slicer.py b/pymc/step_methods/slicer.py
index d8bbc9d2c..0457b3947 100644
--- a/pymc/step_methods/slicer.py
+++ b/pymc/step_methods/slicer.py
@@ -66,34 +66,42 @@ class Slice(ArrayStep):
     def astep(self, q0, logp):
         q0_val = q0.data
         self.w = np.resize(self.w, len(q0_val))  # this is a repmat
-        q = np.copy(q0_val)  # TODO: find out if we need this
+
+        q = np.copy(q0_val)
         ql = np.copy(q0_val)  # l for left boundary
         qr = np.copy(q0_val)  # r for right boudary
-        for i in range(len(q0_val)):
+
+        # The points are not copied, so it's fine to update them inplace in the
+        # loop below
+        q_ra = RaveledVars(q, q0.point_map_info)
+        ql_ra = RaveledVars(ql, q0.point_map_info)
+        qr_ra = RaveledVars(qr, q0.point_map_info)
+
+        for i, wi in enumerate(self.w):
             # uniformly sample from 0 to p(q), but in log space
-            q_ra = RaveledVars(q, q0.point_map_info)
             y = logp(q_ra) - nr.standard_exponential()
-            ql[i] = q[i] - nr.uniform(0, self.w[i])
-            qr[i] = q[i] + self.w[i]
+
+            # Create initial interval
+            ql[i] = q[i] - nr.uniform() * wi  # q[i] + r * w
+            qr[i] = ql[i] + wi  # Equivalent to q[i] + (1-r) * w
+
             # Stepping out procedure
             cnt = 0
-            while y <= logp(
-                RaveledVars(ql, q0.point_map_info)
-            ):  # changed lt to leq  for locally uniform posteriors
-                ql[i] -= self.w[i]
+            while y <= logp(ql_ra):  # changed lt to leq  for locally uniform posteriors
+                ql[i] -= wi
                 cnt += 1
                 if cnt > self.iter_limit:
                     raise RuntimeError(LOOP_ERR_MSG % self.iter_limit)
             cnt = 0
-            while y <= logp(RaveledVars(qr, q0.point_map_info)):
-                qr[i] += self.w[i]
+            while y <= logp(qr_ra):
+                qr[i] += wi
                 cnt += 1
                 if cnt > self.iter_limit:
                     raise RuntimeError(LOOP_ERR_MSG % self.iter_limit)
 
             cnt = 0
             q[i] = nr.uniform(ql[i], qr[i])
-            while logp(q_ra) < y:  # Changed leq to lt, to accomodate for locally flat posteriors
+            while y > logp(q_ra):  # Changed leq to lt, to accomodate for locally flat posteriors
                 # Sample uniformly from slice
                 if q[i] > q0_val[i]:
                     qr[i] = q[i]
@@ -104,18 +112,19 @@ class Slice(ArrayStep):
                 if cnt > self.iter_limit:
                     raise RuntimeError(LOOP_ERR_MSG % self.iter_limit)
 
-            if (
-                self.tune
-            ):  # I was under impression from MacKays lectures that slice width can be tuned without
+            if self.tune:
+                # I was under impression from MacKays lectures that slice width can be tuned without
                 # breaking markovianness. Can we do it regardless of self.tune?(@madanh)
-                self.w[i] = self.w[i] * (self.n_tunes / (self.n_tunes + 1)) + (qr[i] - ql[i]) / (
+                self.w[i] = wi * (self.n_tunes / (self.n_tunes + 1)) + (qr[i] - ql[i]) / (
                     self.n_tunes + 1
-                )  # same as before
-                # unobvious and important: return qr and ql to the same point
-                qr[i] = q[i]
-                ql[i] = q[i]
+                )
+
+            # Set qr and ql to the accepted points (they matter for subsequent iterations)
+            qr[i] = ql[i] = q[i]
+
         if self.tune:
             self.n_tunes += 1
+
         return q
 
     @staticmethod


Test patch:
diff --git a/pymc/tests/test_step.py b/pymc/tests/test_step.py
index 7c89c9100..9df55f57b 100644
--- a/pymc/tests/test_step.py
+++ b/pymc/tests/test_step.py
@@ -60,6 +60,7 @@ from pymc.step_methods import (
 )
 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,
@@ -78,46 +79,58 @@ class TestStepMethods:
         shutil.rmtree(self.temp_dir)
 
     def check_stat(self, check, idata, name):
-        if hasattr(idata, "warmup_posterior"):
-            group = idata.warmup_posterior
-        else:
-            group = idata.posterior
+        group = idata.posterior
         for (var, stat, value, bound) in check:
-            s = stat(group[var].sel(chain=0, draw=slice(2000, None)), axis=0)
-            close_to(s, value, bound)
+            s = stat(group[var].sel(chain=0), axis=0)
+            close_to(s, value, bound, name)
 
-    def test_step_continuous(self):
-        start, model, (mu, C) = mv_simple()
-        unc = np.diag(C) ** 0.5
-        check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
-        _, model_coarse, _ = mv_simple_coarse()
-        with model:
-            steps = (
-                Slice(),
-                HamiltonianMC(scaling=C, is_cov=True, blocked=False),
-                NUTS(scaling=C, is_cov=True, blocked=False),
-                Metropolis(S=C, proposal_dist=MultivariateNormalProposal, blocked=True),
-                Slice(blocked=True),
-                HamiltonianMC(scaling=C, is_cov=True),
-                NUTS(scaling=C, is_cov=True),
-                CompoundStep(
+    @pytest.mark.parametrize(
+        "step_fn, draws",
+        [
+            (lambda C, _: HamiltonianMC(scaling=C, is_cov=True, blocked=False), 1000),
+            (lambda C, _: HamiltonianMC(scaling=C, is_cov=True), 1000),
+            (lambda C, _: NUTS(scaling=C, is_cov=True, blocked=False), 1000),
+            (lambda C, _: NUTS(scaling=C, is_cov=True), 1000),
+            (
+                lambda C, _: CompoundStep(
                     [
                         HamiltonianMC(scaling=C, is_cov=True),
                         HamiltonianMC(scaling=C, is_cov=True, blocked=False),
                     ]
                 ),
-                MLDA(
+                1000,
+            ),
+            # MLDA takes 1/2 of the total test time!
+            (
+                lambda C, model_coarse: MLDA(
                     coarse_models=[model_coarse],
                     base_S=C,
                     base_proposal_dist=MultivariateNormalProposal,
                 ),
-            )
-        for step in steps:
+                1000,
+            ),
+            (lambda *_: Slice(), 2000),
+            (lambda *_: Slice(blocked=True), 2000),
+            (
+                lambda C, _: Metropolis(
+                    S=C, proposal_dist=MultivariateNormalProposal, blocked=True
+                ),
+                4000,
+            ),
+        ],
+        ids=str,
+    )
+    def test_step_continuous(self, step_fn, draws):
+        start, model, (mu, C) = mv_simple()
+        unc = np.diag(C) ** 0.5
+        check = (("x", np.mean, mu, unc / 10), ("x", np.std, unc, unc / 10))
+        _, model_coarse, _ = mv_simple_coarse()
+        with model:
+            step = step_fn(C, model_coarse)
             idata = sample(
-                0,
-                tune=8000,
+                tune=1000,
+                draws=draws,
                 chains=1,
-                discard_tuned_samples=False,
                 step=step,
                 start=start,
                 model=model,
@@ -126,49 +139,58 @@ class TestStepMethods:
             self.check_stat(check, idata, step.__class__.__name__)
 
     def test_step_discrete(self):
-        if aesara.config.floatX == "float32":
-            return  # Cannot use @skip because it only skips one iteration of the yield
         start, model, (mu, C) = mv_simple_discrete()
         unc = np.diag(C) ** 0.5
         check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
         with model:
-            steps = (Metropolis(S=C, proposal_dist=MultivariateNormalProposal),)
-        for step in steps:
+            step = Metropolis(S=C, proposal_dist=MultivariateNormalProposal)
             idata = sample(
-                20000, tune=0, step=step, start=start, model=model, random_seed=1, chains=1
+                tune=1000,
+                draws=2000,
+                chains=1,
+                step=step,
+                start=start,
+                model=model,
+                random_seed=1,
             )
             self.check_stat(check, idata, step.__class__.__name__)
 
-    def test_step_categorical(self):
+    @pytest.mark.parametrize("proposal", ["uniform", "proportional"])
+    def test_step_categorical(self, proposal):
         start, model, (mu, C) = simple_categorical()
         unc = C**0.5
         check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
         with model:
-            steps = (
-                CategoricalGibbsMetropolis([model.x], proposal="uniform"),
-                CategoricalGibbsMetropolis([model.x], proposal="proportional"),
+            step = CategoricalGibbsMetropolis([model.x], proposal=proposal)
+            idata = sample(
+                tune=1000,
+                draws=2000,
+                chains=1,
+                step=step,
+                start=start,
+                model=model,
+                random_seed=1,
             )
-        for step in steps:
-            idata = sample(8000, tune=0, step=step, start=start, model=model, random_seed=1)
             self.check_stat(check, idata, step.__class__.__name__)
 
 
 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)
@@ -182,26 +204,32 @@ class TestMetropolisProposal:
 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)
+
+    def test_name(self):
+        with Model() as m:
+            c1 = HalfNormal("c1")
+            c2 = HalfNormal("c2")
+
+            step1 = NUTS([c1])
+            step2 = Slice([c2])
+            step = CompoundStep([step1, step2])
+        assert step.name == "Compound[nuts, slice]"
 
 
 class TestAssignStepMethods:
@@ -209,32 +237,37 @@ class TestAssignStepMethods:
         """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):
@@ -254,7 +287,8 @@ class TestAssignStepMethods:
             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):
@@ -266,7 +300,8 @@ class TestAssignStepMethods:
 
         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
@@ -274,7 +309,8 @@ class TestAssignStepMethods:
 
         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)
 
 
@@ -1306,7 +1342,8 @@ class TestRVsAssignmentSteps:
             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
             )
@@ -1323,7 +1360,8 @@ class TestRVsAssignmentSteps:
             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
             )
@@ -1333,7 +1371,8 @@ class TestRVsAssignmentSteps:
             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)