Repository: pymc-devs/pymc

Issue:
Add return types in docstrings for logprob.py
Add return types in docstring for all methods in `logprob.py` following the numpy style. Theyre missing and it would be nice to include them for consistency and readability

https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_numpy.html
https://github.com/pymc-devs/pymc3/blob/main/pymc3/distributions/logprob.py#L48


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5819?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5819](https://codecov.io/gh/pymc-devs/pymc/pull/5819?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (2db70ce) 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 `22.41%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5819/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/5819?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    #5819       +/-   ##
===========================================
- Coverage   89.39%   66.97%   -22.42%     
===========================================
  Files          74       74               
  Lines       13767    13767               
===========================================
- Hits        12307     9221     -3086     
- Misses       1460     4546     +3086     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/distributions/logprob.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2xvZ3Byb2IucHk=) | `86.71% <100.00%> (-10.94%)` | :arrow_down: |
| [pymc/step\_methods/mlda.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvbWxkYS5weQ==) | `12.53% <0.00%> (-83.85%)` | :arrow_down: |
| [pymc/tuning/starting.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/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/step\_methods/slicer.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvc2xpY2VyLnB5) | `26.76% <0.00%> (-69.02%)` | :arrow_down: |
| [pymc/distributions/mixture.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/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/func\_utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/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: |
| [pymc/printing.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9wcmludGluZy5weQ==) | `21.78% <0.00%> (-64.36%)` | :arrow_down: |
| [pymc/step\_methods/hmc/hmc.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvaG1jL2htYy5weQ==) | `33.96% <0.00%> (-58.50%)` | :arrow_down: |
| [pymc/distributions/bound.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2JvdW5kLnB5) | `45.54% <0.00%> (-54.46%)` | :arrow_down: |
| [pymc/distributions/discrete.py](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2Rpc2NyZXRlLnB5) | `49.20% <0.00%> (-50.54%)` | :arrow_down: |
| ... and [32 more](https://codecov.io/gh/pymc-devs/pymc/pull/5819/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |


Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 312782f4c..62bbc1249 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -1,7 +1,15 @@
 # Guidelines for Contributing
 
-The guidelines for contributing to the PyMC project are available
-on [our documentation](https://docs.pymc.io/en/latest/contributing/index.html)
+Thank you for being interested in contributing to PyMC. PyMC is an open source, collective effort, and everyone is welcome to contribute. There are many ways in which you can help make it better. Please check the latest information for contributing to the PyMC project on [this guidelines](https://docs.pymc.io/en/latest/contributing/index.html).
+
+Quick links
+-----------
+
+* [Pull request (PR) step-by-step ](https://docs.pymc.io/en/latest/contributing/pr_tutorial.html)
+* [Pull request (PR) checklist](https://docs.pymc.io/en/latest/contributing/pr_checklist.html)
+* [Python style guide with pre-commit](https://docs.pymc.io/en/latest/contributing/python_style.html)
+* [Running the test suite](https://docs.pymc.io/en/latest/contributing/running_the_test_suite.html)
+* [Submitting a bug report or feature request](https://github.com/pymc-devs/pymc/issues)
 
 <!-- Commented out because our Docker image is outdated/broken.
 ## Developing in Docker
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/aesaraf.py b/pymc/aesaraf.py
index 0160635d0..a3d19659c 100644
--- a/pymc/aesaraf.py
+++ b/pymc/aesaraf.py
@@ -52,6 +52,10 @@ from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.scalar.basic import Cast
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.random.var import (
+    RandomGeneratorSharedVariable,
+    RandomStateSharedVariable,
+)
 from aesara.tensor.shape import SpecifyShape
 from aesara.tensor.sharedvar import SharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
@@ -60,9 +64,7 @@ from aesara.tensor.var import TensorConstant, TensorVariable
 from pymc.exceptions import ShapeError
 from pymc.vartypes import continuous_types, isgenerator, typefilter
 
-PotentialShapeType = Union[
-    int, np.ndarray, Tuple[Union[int, Variable], ...], List[Union[int, Variable]], Variable
-]
+PotentialShapeType = Union[int, np.ndarray, Sequence[Union[int, Variable]], TensorVariable]
 
 
 __all__ = [
@@ -165,6 +167,7 @@ def change_rv_size(
         new_size = (new_size,)
 
     # Extract the RV node that is to be resized, together with its inputs, name and tag
+    assert rv.owner.op is not None
     if isinstance(rv.owner.op, SpecifyShape):
         rv = rv.owner.inputs[0]
     rv_node = rv.owner
@@ -894,18 +897,14 @@ aesara.compile.optdb["canonicalize"].register(
 )
 
 
-def find_rng_nodes(variables: Iterable[TensorVariable]):
+def find_rng_nodes(
+    variables: Iterable[Variable],
+) -> List[Union[RandomStateSharedVariable, RandomGeneratorSharedVariable]]:
     """Return RNG variables in a graph"""
     return [
         node
         for node in graph_inputs(variables)
-        if isinstance(
-            node,
-            (
-                at.random.var.RandomStateSharedVariable,
-                at.random.var.RandomGeneratorSharedVariable,
-            ),
-        )
+        if isinstance(node, (RandomStateSharedVariable, RandomGeneratorSharedVariable))
     ]
 
 
@@ -921,6 +920,7 @@ def reseed_rngs(
         np.random.PCG64(sub_seed) for sub_seed in np.random.SeedSequence(seed).spawn(len(rngs))
     ]
     for rng, bit_generator in zip(rngs, bit_generators):
+        new_rng: Union[np.random.RandomState, np.random.Generator]
         if isinstance(rng, at.random.var.RandomStateSharedVariable):
             new_rng = np.random.RandomState(bit_generator)
         else:
@@ -980,6 +980,9 @@ def compile_pymc(
         and isinstance(var.owner.op, (RandomVariable, MeasurableVariable))
         and var not in inputs
     ):
+        # All nodes in `vars_between(inputs, outputs)` have owners.
+        # But mypy doesn't know, so we just assert it:
+        assert random_var.owner.op is not None
         if isinstance(random_var.owner.op, RandomVariable):
             rng = random_var.owner.inputs[0]
             if not hasattr(rng, "default_update"):
diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
index 35fce9e22..22d766257 100644
--- a/pymc/distributions/distribution.py
+++ b/pymc/distributions/distribution.py
@@ -198,8 +198,8 @@ class Distribution(metaclass=DistributionMeta):
         total_size=None,
         transform=UNSET,
         **kwargs,
-    ) -> RandomVariable:
-        """Adds a RandomVariable corresponding to a PyMC distribution to the current model.
+    ) -> TensorVariable:
+        """Adds a tensor variable corresponding to a PyMC distribution to the current model.
 
         Note that all remaining kwargs must be compatible with ``.dist()``
 
@@ -231,8 +231,8 @@ class Distribution(metaclass=DistributionMeta):
 
         Returns
         -------
-        rv : RandomVariable
-            The created RV, registered in the Model.
+        rv : TensorVariable
+            The created random variable tensor, registered in the Model.
         """
 
         try:
@@ -296,8 +296,8 @@ class Distribution(metaclass=DistributionMeta):
         *,
         shape: Optional[Shape] = None,
         **kwargs,
-    ) -> RandomVariable:
-        """Creates a RandomVariable corresponding to the `cls` distribution.
+    ) -> TensorVariable:
+        """Creates a tensor variable corresponding to the `cls` distribution.
 
         Parameters
         ----------
@@ -314,8 +314,8 @@ class Distribution(metaclass=DistributionMeta):
 
         Returns
         -------
-        rv : RandomVariable
-            The created RV.
+        rv : TensorVariable
+            The created random variable tensor.
         """
         if "testval" in kwargs:
             kwargs.pop("testval")
@@ -653,8 +653,8 @@ class DensityDist(NoDistribution):
         name : str
         dist_params : Tuple
             A sequence of the distribution's parameter. These will be converted into
-            Aesara tensors internally. These parameters could be other ``RandomVariable``
-            instances.
+            Aesara tensors internally. These parameters could be other ``TensorVariable``
+            instances created from , optionally created via ``RandomVariable`` ``Op``s.
         logp : Optional[Callable]
             A callable that calculates the log density of some given observed ``value``
             conditioned on certain distribution parameter values. It must have the
diff --git a/pymc/distributions/logprob.py b/pymc/distributions/logprob.py
index 9b92f8acf..fb2b041ff 100644
--- a/pymc/distributions/logprob.py
+++ b/pymc/distributions/logprob.py
@@ -39,7 +39,9 @@ from aesara.tensor.var import TensorVariable
 from pymc.aesaraf import floatX
 
 
-def _get_scaling(total_size: Optional[Union[int, Sequence[int]]], shape, ndim: int):
+def _get_scaling(
+    total_size: Optional[Union[int, Sequence[int]]], shape, ndim: int
+) -> TensorVariable:
     """
     Gets scaling constant for logp.
 
@@ -288,14 +290,14 @@ def logp(rv: TensorVariable, value) -> TensorVariable:
             raise NotImplementedError("PyMC could not infer logp of input variable.") from exc
 
 
-def logcdf(rv, value):
+def logcdf(rv: TensorVariable, value) -> TensorVariable:
     """Return the log-cdf graph of a Random Variable"""
 
     value = at.as_tensor_variable(value, dtype=rv.dtype)
     return logcdf_aeppl(rv, value)
 
 
-def ignore_logprob(rv):
+def ignore_logprob(rv: TensorVariable) -> TensorVariable:
     """Return a duplicated variable that is ignored when creating Aeppl logprob graphs
 
     This is used in SymbolicDistributions that use other RVs as inputs but account
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)