pymc-extras 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

tests/distributions/__init__.py

@@ -1,19 +0,0 @@
-#   Copyright 2022 The PyMC Developers
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-
-
-from pymc_extras.distributions import histogram_utils
-from pymc_extras.distributions.histogram_utils import histogram_approximation
-
-__all__ = ["histogram_utils", "histogram_approximation"]

tests/distributions/test_continuous.py

@@ -1,185 +0,0 @@
-#   Copyright 2020 The PyMC Developers
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-import numpy as np
-import pymc as pm
-
-# general imports
-import pytest
-import scipy.stats.distributions as sp
-
-
-# test support imports from pymc
-from pymc.testing import (
-    BaseTestDistributionRandom,
-    Domain,
-    R,
-    Rplus,
-    Rplusbig,
-    assert_support_point_is_expected,
-    check_logcdf,
-    check_logp,
-    seeded_scipy_distribution_builder,
-    select_by_precision,
-)
-
-# the distributions to be tested
-from pymc_extras.distributions import Chi, GenExtreme, Maxwell
-
-
-class TestGenExtremeClass:
-    """
-    Wrapper class so that tests of experimental additions can be dropped into
-    PyMC directly on adoption.
-
-    pm.logp(GenExtreme.dist(mu=0.,sigma=1.,xi=0.5),value=-0.01)
-    """
-
-    def test_logp(self):
-        def ref_logp(value, mu, sigma, xi):
-            if 1 + xi * (value - mu) / sigma > 0:
-                return sp.genextreme.logpdf(value, c=-xi, loc=mu, scale=sigma)
-            else:
-                return -np.inf
-
-        check_logp(
-            GenExtreme,
-            R,
-            {
-                "mu": R,
-                "sigma": Rplusbig,
-                "xi": Domain([-1, -0.99, -0.5, 0, 0.5, 0.99, 1]),
-            },
-            ref_logp,
-        )
-
-    def test_logcdf(self):
-        def ref_logcdf(value, mu, sigma, xi):
-            if 1 + xi * (value - mu) / sigma > 0:
-                return sp.genextreme.logcdf(value, c=-xi, loc=mu, scale=sigma)
-            else:
-                return -np.inf
-
-        check_logcdf(
-            GenExtreme,
-            R,
-            {
-                "mu": R,
-                "sigma": Rplusbig,
-                "xi": Domain([-1, -0.99, -0.5, 0, 0.5, 0.99, 1]),
-            },
-            ref_logcdf,
-            decimal=select_by_precision(float64=6, float32=2),
-        )
-
-    @pytest.mark.parametrize(
-        "mu, sigma, xi, size, expected",
-        [
-            (0, 1, 0, None, 0),
-            (1, np.arange(1, 4), 0.1, None, 1 + np.arange(1, 4) * (1.1**-0.1 - 1) / 0.1),
-            (np.arange(5), 1, 0.1, None, np.arange(5) + (1.1**-0.1 - 1) / 0.1),
-            (
-                0,
-                1,
-                np.linspace(-0.2, 0.2, 6),
-                None,
-                ((1 + np.linspace(-0.2, 0.2, 6)) ** -np.linspace(-0.2, 0.2, 6) - 1)
-                / np.linspace(-0.2, 0.2, 6),
-            ),
-            (1, 2, 0.1, 5, np.full(5, 1 + 2 * (1.1**-0.1 - 1) / 0.1)),
-            (
-                np.arange(6),
-                np.arange(1, 7),
-                np.linspace(-0.2, 0.2, 6),
-                (3, 6),
-                np.full(
-                    (3, 6),
-                    np.arange(6)
-                    + np.arange(1, 7)
-                    * ((1 + np.linspace(-0.2, 0.2, 6)) ** -np.linspace(-0.2, 0.2, 6) - 1)
-                    / np.linspace(-0.2, 0.2, 6),
-                ),
-            ),
-        ],
-    )
-    def test_genextreme_support_point(self, mu, sigma, xi, size, expected):
-        with pm.Model() as model:
-            GenExtreme("x", mu=mu, sigma=sigma, xi=xi, size=size)
-        assert_support_point_is_expected(model, expected)
-
-    def test_gen_extreme_scipy_kwarg(self):
-        dist = GenExtreme.dist(xi=1, scipy=False)
-        assert dist.owner.inputs[-1].eval() == 1
-
-        dist = GenExtreme.dist(xi=1, scipy=True)
-        assert dist.owner.inputs[-1].eval() == -1
-
-
-class TestGenExtreme(BaseTestDistributionRandom):
-    pymc_dist = GenExtreme
-    pymc_dist_params = {"mu": 0, "sigma": 1, "xi": -0.1}
-    expected_rv_op_params = {"mu": 0, "sigma": 1, "xi": -0.1}
-    # Notice, using different parametrization of xi sign to scipy
-    reference_dist_params = {"loc": 0, "scale": 1, "c": 0.1}
-    reference_dist = seeded_scipy_distribution_builder("genextreme")
-    tests_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_pymc_draws_match_reference",
-        "check_rv_size",
-    ]
-
-
-class TestChiClass:
-    """
-    Wrapper class so that tests of experimental additions can be dropped into
-    PyMC directly on adoption.
-    """
-
-    def test_logp(self):
-        check_logp(
-            Chi,
-            Rplus,
-            {"nu": Rplus},
-            lambda value, nu: sp.chi.logpdf(value, df=nu),
-        )
-
-    def test_logcdf(self):
-        check_logcdf(
-            Chi,
-            Rplus,
-            {"nu": Rplus},
-            lambda value, nu: sp.chi.logcdf(value, df=nu),
-        )
-
-
-class TestMaxwell:
-    """
-    Wrapper class so that tests of experimental additions can be dropped into
-    PyMC directly on adoption.
-    """
-
-    def test_logp(self):
-        check_logp(
-            Maxwell,
-            Rplus,
-            {"a": Rplus},
-            lambda value, a: sp.maxwell.logpdf(value, scale=a),
-        )
-
-    def test_logcdf(self):
-        check_logcdf(
-            Maxwell,
-            Rplus,
-            {"a": Rplus},
-            lambda value, a: sp.maxwell.logcdf(value, scale=a),
-        )

tests/distributions/test_discrete.py

@@ -1,210 +0,0 @@
-#   Copyright 2023 The PyMC Developers
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-import numpy as np
-import pymc as pm
-import pytensor
-import pytensor.tensor as pt
-import pytest
-import scipy.stats
-
-from pymc.logprob.utils import ParameterValueError
-from pymc.testing import (
-    BaseTestDistributionRandom,
-    Domain,
-    I,
-    Rplus,
-    assert_support_point_is_expected,
-    check_logp,
-    discrete_random_tester,
-)
-from pytensor import config
-
-from pymc_extras.distributions import (
-    BetaNegativeBinomial,
-    GeneralizedPoisson,
-    Skellam,
-)
-
-
-class TestGeneralizedPoisson:
-    class TestRandomVariable(BaseTestDistributionRandom):
-        pymc_dist = GeneralizedPoisson
-        pymc_dist_params = {"mu": 4.0, "lam": 1.0}
-        expected_rv_op_params = {"mu": 4.0, "lam": 1.0}
-        tests_to_run = [
-            "check_pymc_params_match_rv_op",
-            "check_rv_size",
-        ]
-
-        def test_random_matches_poisson(self):
-            discrete_random_tester(
-                dist=self.pymc_dist,
-                paramdomains={"mu": Rplus, "lam": Domain([0], edges=(None, None))},
-                ref_rand=lambda mu, lam, size: scipy.stats.poisson.rvs(mu, size=size),
-            )
-
-        @pytest.mark.parametrize("mu", (2.5, 20, 50))
-        def test_random_lam_expected_moments(self, mu):
-            lam = np.array([-0.9, -0.7, -0.2, 0, 0.2, 0.7, 0.9])
-            dist = self.pymc_dist.dist(mu=mu, lam=lam, size=(10_000, len(lam)))
-            draws = dist.eval()
-
-            expected_mean = mu / (1 - lam)
-            np.testing.assert_allclose(draws.mean(0), expected_mean, rtol=1e-1)
-
-            expected_std = np.sqrt(mu / (1 - lam) ** 3)
-            np.testing.assert_allclose(draws.std(0), expected_std, rtol=1e-1)
-
-    def test_logp_matches_poisson(self):
-        # We are only checking this distribution for lambda=0 where it's equivalent to Poisson.
-        mu = pt.scalar("mu")
-        lam = pt.scalar("lam")
-        value = pt.vector("value", dtype="int64")
-
-        logp = pm.logp(GeneralizedPoisson.dist(mu, lam), value)
-        logp_fn = pytensor.function([value, mu, lam], logp)
-
-        test_value = np.array([0, 1, 2, 30])
-        for test_mu in (0.01, 0.1, 0.9, 1, 1.5, 20, 100):
-            np.testing.assert_allclose(
-                logp_fn(test_value, test_mu, lam=0),
-                scipy.stats.poisson.logpmf(test_value, test_mu),
-                rtol=1e-7 if config.floatX == "float64" else 1e-5,
-            )
-
-        # Check out-of-bounds values
-        value = pt.scalar("value")
-        logp = pm.logp(GeneralizedPoisson.dist(mu, lam), value)
-        logp_fn = pytensor.function([value, mu, lam], logp)
-
-        logp_fn(-1, mu=5, lam=0) == -np.inf
-        logp_fn(9, mu=5, lam=-1) == -np.inf
-
-        # Check mu/lam restrictions
-        with pytest.raises(ParameterValueError):
-            logp_fn(1, mu=1, lam=2)
-
-        with pytest.raises(ParameterValueError):
-            logp_fn(1, mu=0, lam=0)
-
-        with pytest.raises(ParameterValueError):
-            logp_fn(1, mu=1, lam=-1)
-
-    def test_logp_lam_expected_moments(self):
-        mu = 30
-        lam = np.array([-0.9, -0.7, -0.2, 0, 0.2, 0.7, 0.9])
-        with pm.Model():
-            x = GeneralizedPoisson("x", mu=mu, lam=lam)
-            trace = pm.sample(chains=1, draws=10_000, random_seed=96).posterior
-
-        expected_mean = mu / (1 - lam)
-        np.testing.assert_allclose(trace["x"].mean(("chain", "draw")), expected_mean, rtol=1e-1)
-
-        expected_std = np.sqrt(mu / (1 - lam) ** 3)
-        np.testing.assert_allclose(trace["x"].std(("chain", "draw")), expected_std, rtol=1e-1)
-
-    @pytest.mark.parametrize(
-        "mu, lam, size, expected",
-        [
-            (50, [-0.6, 0, 0.6], None, np.floor(50 / (1 - np.array([-0.6, 0, 0.6])))),
-            ([5, 50], -0.1, (4, 2), np.full((4, 2), np.floor(np.array([5, 50]) / 1.1))),
-        ],
-    )
-    def test_moment(self, mu, lam, size, expected):
-        with pm.Model() as model:
-            GeneralizedPoisson("x", mu=mu, lam=lam, size=size)
-        assert_support_point_is_expected(model, expected)
-
-
-class TestBetaNegativeBinomial:
-    """
-    Wrapper class so that tests of experimental additions can be dropped into
-    PyMC directly on adoption.
-    """
-
-    def test_logp(self):
-        """
-
-        Beta Negative Binomial logp function test values taken from R package as
-        there is currently no implementation in scipy.
-        https://github.com/scipy/scipy/issues/17330
-
-        The test values can be generated in R with the following code:
-
-        .. code-block:: r
-
-            library(extraDistr)
-
-            create.test.rows <- function(alpha, beta, r, x) {
-                logp <- dbnbinom(x, alpha, beta, r, log=TRUE)
-                paste0("(", paste(alpha, beta, r, x, logp, sep=", "), ")")
-            }
-
-            x <- c(0, 1, 250, 5000)
-            print(create.test.rows(1, 1, 1, x), quote=FALSE)
-            print(create.test.rows(1, 1, 10, x), quote=FALSE)
-            print(create.test.rows(1, 10, 1, x), quote=FALSE)
-            print(create.test.rows(10, 1, 1, x), quote=FALSE)
-            print(create.test.rows(10, 10, 10, x), quote=FALSE)
-
-        """
-        alpha, beta, r, value = pt.scalars("alpha", "beta", "r", "value")
-        logp = pm.logp(BetaNegativeBinomial.dist(alpha, beta, r), value)
-        logp_fn = pytensor.function([value, alpha, beta, r], logp)
-
-        tests = [
-            # 1, 1, 1
-            (1, 1, 1, 0, -0.693147180559945),
-            (1, 1, 1, 1, -1.79175946922805),
-            (1, 1, 1, 250, -11.0548820266432),
-            (1, 1, 1, 5000, -17.0349862828565),
-            # 1, 1, 10
-            (1, 1, 10, 0, -2.39789527279837),
-            (1, 1, 10, 1, -2.58021682959232),
-            (1, 1, 10, 250, -8.82261694534392),
-            (1, 1, 10, 5000, -14.7359968760473),
-            # 1, 10, 1
-            (1, 10, 1, 0, -2.39789527279837),
-            (1, 10, 1, 1, -2.58021682959232),
-            (1, 10, 1, 250, -8.82261694534418),
-            (1, 10, 1, 5000, -14.7359968760446),
-            # 10, 1, 1
-            (10, 1, 1, 0, -0.0953101798043248),
-            (10, 1, 1, 1, -2.58021682959232),
-            (10, 1, 1, 250, -43.5891148758123),
-            (10, 1, 1, 5000, -76.2953173311091),
-            # 10, 10, 10
-            (10, 10, 10, 0, -5.37909807285049),
-            (10, 10, 10, 1, -4.17512526852455),
-            (10, 10, 10, 250, -21.781591505836),
-            (10, 10, 10, 5000, -53.4836799634603),
-        ]
-        for test_alpha, test_beta, test_r, test_value, expected_logp in tests:
-            np.testing.assert_allclose(
-                logp_fn(test_value, test_alpha, test_beta, test_r), expected_logp
-            )
-
-
-class TestSkellam:
-    def test_logp(self):
-        # Scipy Skellam underflows to -inf earlier than PyMC
-        Rplus_small = Domain([0, 0.01, 0.1, 0.9, 0.99, 1, 1.5, 2, 10, np.inf])
-        # Suppress warnings coming from Scipy logpmf implementation
-        with np.errstate(divide="ignore", invalid="ignore"):
-            check_logp(
-                Skellam,
-                I,
-                {"mu1": Rplus_small, "mu2": Rplus_small},
-                lambda value, mu1, mu2: scipy.stats.skellam.logpmf(value, mu1, mu2),
-            )

tests/distributions/test_discrete_markov_chain.py

@@ -1,258 +0,0 @@
-import numpy as np
-import pymc as pm
-
-# general imports
-import pytensor.tensor as pt
-import pytest
-
-from pymc.distributions import Categorical
-from pymc.distributions.shape_utils import change_dist_size
-from pymc.logprob.utils import ParameterValueError
-from pymc.sampling.mcmc import assign_step_methods
-
-from pymc_extras.distributions.timeseries import (
-    DiscreteMarkovChain,
-    DiscreteMarkovChainGibbsMetropolis,
-)
-
-
-def transition_probability_tests(steps, n_states, n_lags, n_draws, atol):
-    P = np.full((n_states,) * (n_lags + 1), 1 / n_states)
-    x0 = pm.Categorical.dist(p=np.ones(n_states) / n_states)
-
-    chain = DiscreteMarkovChain.dist(
-        P=pt.as_tensor_variable(P), init_dist=x0, steps=steps, n_lags=n_lags
-    )
-
-    draws = pm.draw(chain, n_draws, random_seed=172)
-
-    # Test x0 is uniform over n_states
-    for i in range(n_lags):
-        assert np.allclose(
-            np.histogram(draws[:, ..., i], bins=n_states)[0] / n_draws, 1 / n_states, atol=atol
-        )
-
-    n_grams = [[tuple(row[i : i + n_lags + 1]) for i in range(len(row) - n_lags)] for row in draws]
-    freq_table = np.zeros((n_states,) * (n_lags + 1))
-
-    for row in n_grams:
-        for ngram in row:
-            freq_table[ngram] += 1
-    freq_table /= freq_table.sum(axis=-1)[:, None]
-
-    # Test continuation probabilities match P
-    assert np.allclose(P, freq_table, atol=atol)
-
-
-class TestDiscreteMarkovRV:
-    def test_fail_if_P_not_square(self):
-        P = pt.eye(3, 2)
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
-        with pytest.raises(ParameterValueError):
-            pm.logp(chain, np.zeros((3,))).eval()
-
-    def test_fail_if_P_not_valid(self):
-        P = pt.zeros((3, 3))
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
-        with pytest.raises(ParameterValueError):
-            pm.logp(chain, np.zeros((3,))).eval()
-
-    def test_high_dimensional_P(self):
-        P = pm.Dirichlet.dist(a=pt.ones(3), size=(3, 3, 3))
-        n_lags = 3
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-        chain = DiscreteMarkovChain.dist(P=P, steps=10, init_dist=x0, n_lags=n_lags)
-        draws = pm.draw(chain, 10)
-        logp = pm.logp(chain, draws)
-
-    def test_default_init_dist_warns_user(self):
-        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
-
-        with pytest.warns(UserWarning):
-            DiscreteMarkovChain.dist(P=P, steps=3)
-
-    def test_logp_shape(self):
-        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-        # Test with steps
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
-        draws = pm.draw(chain, 5)
-        logp = pm.logp(chain, draws).eval()
-
-        assert logp.shape == (5,)
-
-        # Test with shape
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(3,))
-        draws = pm.draw(chain, 5)
-        logp = pm.logp(chain, draws).eval()
-
-        assert logp.shape == (5,)
-
-    def test_logp_with_default_init_dist(self):
-        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-        value = np.array([0, 1, 2])
-        logp_expected = np.log((1 / 3) * 0.5 * 0.3)
-
-        # Test dist directly
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
-        logp_eval = pm.logp(chain, value).eval()
-        np.testing.assert_allclose(logp_eval, logp_expected, rtol=1e-6)
-
-        # Test via Model
-        with pm.Model() as m:
-            DiscreteMarkovChain("chain", P=P, init_dist=x0, steps=3)
-        model_logp_eval = m.compile_logp()({"chain": value})
-        np.testing.assert_allclose(model_logp_eval, logp_expected, rtol=1e-6)
-
-    def test_logp_with_user_defined_init_dist(self):
-        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
-        x0 = pm.Categorical.dist(p=[0.2, 0.6, 0.2])
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
-
-        logp = pm.logp(chain, [0, 1, 2]).eval()
-        assert logp == np.log(0.2 * 0.5 * 0.3)
-
-    def test_moment_function(self):
-        P_np = np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]])
-
-        x0_np = np.array([0, 1, 0])
-
-        P = pt.as_tensor_variable(P_np)
-        x0 = pm.Categorical.dist(p=x0_np.tolist())
-        n_steps = 3
-
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=n_steps)
-
-        chain_np = np.empty(shape=n_steps + 1, dtype="int8")
-        chain_np[0] = np.argmax(x0_np)
-        for i in range(n_steps):
-            state = chain_np[i]
-            chain_np[i + 1] = np.argmax(P_np[state])
-
-        dmc_chain = pm.distributions.distribution.support_point(chain).eval()
-
-        assert np.allclose(dmc_chain, chain_np)
-
-    def test_define_steps_via_shape_arg(self):
-        P = pt.full((3, 3), 1 / 3)
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(3,))
-        assert chain.eval().shape == (3,)
-
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(3, 2))
-        assert chain.eval().shape == (3, 2)
-
-    def test_define_steps_via_dim_arg(self):
-        coords = {"steps": [1, 2, 3]}
-
-        with pm.Model(coords=coords):
-            P = pt.full((3, 3), 1 / 3)
-            x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-            chain = DiscreteMarkovChain("chain", P=P, init_dist=x0, dims=["steps"])
-
-        assert chain.eval().shape == (3,)
-
-    def test_dims_when_steps_are_defined(self):
-        coords = {"steps": [1, 2, 3, 4]}
-
-        with pm.Model(coords=coords):
-            P = pt.full((3, 3), 1 / 3)
-            x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-            chain = DiscreteMarkovChain("chain", P=P, steps=3, init_dist=x0, dims=["steps"])
-
-        assert chain.eval().shape == (4,)
-
-    def test_multiple_dims_with_steps(self):
-        coords = {"steps": [1, 2, 3], "mc_chains": [1, 2, 3]}
-
-        with pm.Model(coords=coords):
-            P = pt.full((3, 3), 1 / 3)
-            x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-            chain = DiscreteMarkovChain(
-                "chain", P=P, steps=2, init_dist=x0, dims=["steps", "mc_chains"]
-            )
-
-        assert chain.eval().shape == (3, 3)
-
-    def test_mutiple_dims_with_steps_and_init_dist(self):
-        coords = {"steps": [1, 2, 3], "mc_chains": [1, 2, 3]}
-
-        with pm.Model(coords=coords):
-            P = pt.full((3, 3), 1 / 3)
-            x0 = pm.Categorical.dist(p=[0.1, 0.1, 0.8], size=(3,))
-            chain = DiscreteMarkovChain(
-                "chain", P=P, init_dist=x0, steps=2, dims=["steps", "mc_chains"]
-            )
-
-        assert chain.eval().shape == (3, 3)
-
-    def test_multiple_lags_with_data(self):
-        with pm.Model():
-            P = pt.full((3, 3, 3), 1 / 3)
-            x0 = pm.Categorical.dist(p=[0.1, 0.1, 0.8], size=2)
-            data = pm.draw(x0, 100)
-
-            chain = DiscreteMarkovChain("chain", P=P, init_dist=x0, n_lags=2, observed=data)
-
-        assert chain.eval().shape == (100, 2)
-
-    def test_random_draws(self):
-        transition_probability_tests(steps=3, n_states=2, n_lags=1, n_draws=2500, atol=0.05)
-        transition_probability_tests(steps=3, n_states=2, n_lags=3, n_draws=7500, atol=0.05)
-
-    def test_change_size_univariate(self):
-        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
-        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
-
-        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(100, 5))
-
-        new_rw = change_dist_size(chain, new_size=(7,))
-        assert tuple(new_rw.shape.eval()) == (7, 5)
-
-        new_rw = change_dist_size(chain, new_size=(4, 3), expand=True)
-        assert tuple(new_rw.shape.eval()) == (4, 3, 100, 5)
-
-    def test_mcmc_sampling(self):
-        with pm.Model(coords={"step": range(100)}) as model:
-            init_dist = Categorical.dist(p=[0.5, 0.5])
-            markov_chain = DiscreteMarkovChain(
-                "markov_chain",
-                P=[[0.1, 0.9], [0.1, 0.9]],
-                init_dist=init_dist,
-                shape=(100,),
-                dims="step",
-            )
-
-            _, assigned_step_methods = assign_step_methods(model)
-            assert assigned_step_methods[DiscreteMarkovChainGibbsMetropolis] == [
-                model.rvs_to_values[markov_chain]
-            ]
-
-            # Sampler needs no tuning
-            idata = pm.sample(
-                tune=0, chains=4, draws=250, progressbar=False, compute_convergence_checks=False
-            )
-
-        np.testing.assert_allclose(
-            idata.posterior["markov_chain"].isel(step=0).mean(("chain", "draw")),
-            0.5,
-            atol=0.05,
-        )
-
-        np.testing.assert_allclose(
-            idata.posterior["markov_chain"].isel(step=slice(1, None)).mean(("chain", "draw")),
-            0.9,
-            atol=0.05,
-        )
-
-        assert pm.stats.ess(idata, method="tail").min() > 950