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

tests/distributions/test_multivariate.py

@@ -1,304 +0,0 @@
-import numpy as np
-import pymc as pm
-import pytensor
-import pytest
-
-import pymc_extras as pmx
-
-
-class TestR2D2M2CP:
-    @pytest.fixture(autouse=True)
-    def fast_compile(self):
-        with pytensor.config.change_flags(mode="FAST_COMPILE", exception_verbosity="high"):
-            yield
-
-    @pytest.fixture(autouse=True)
-    def model(self):
-        # every method is within a model
-        with pm.Model() as model:
-            yield model
-
-    @pytest.fixture(params=[True, False], ids=["centered", "non-centered"])
-    def centered(self, request):
-        return request.param
-
-    @pytest.fixture(params=[["a"], ["a", "b"], ["one"]])
-    def dims(self, model: pm.Model, request):
-        for i, c in enumerate(request.param):
-            if c == "one":
-                model.add_coord(c, range(1))
-            else:
-                model.add_coord(c, range((i + 2) ** 2))
-        return request.param
-
-    @pytest.fixture
-    def input_shape(self, dims, model):
-        return [int(model.dim_lengths[d].eval()) for d in dims]
-
-    @pytest.fixture
-    def output_shape(self, dims, model):
-        *hierarchy, _ = dims
-        return [int(model.dim_lengths[d].eval()) for d in hierarchy]
-
-    @pytest.fixture
-    def input_std(self, input_shape):
-        return np.ones(input_shape)
-
-    @pytest.fixture
-    def output_std(self, output_shape):
-        return np.ones(output_shape)
-
-    @pytest.fixture
-    def r2(self):
-        return 0.8
-
-    @pytest.fixture(params=[None, 0.1], ids=["r2-std", "no-r2-std"])
-    def r2_std(self, request):
-        return request.param
-
-    @pytest.fixture(params=["true", "false", "limit-1", "limit-0", "limit-all"])
-    def positive_probs(self, input_std, request):
-        if request.param == "true":
-            return np.full_like(input_std, 0.5)
-        elif request.param == "false":
-            return 0.5
-        elif request.param == "limit-1":
-            ret = np.full_like(input_std, 0.5)
-            ret[..., 0] = 1
-            return ret
-        elif request.param == "limit-0":
-            ret = np.full_like(input_std, 0.5)
-            ret[..., 0] = 0
-            return ret
-        elif request.param == "limit-all":
-            return np.full_like(input_std, 0)
-
-    @pytest.fixture(params=[True, False], ids=["probs-std", "no-probs-std"])
-    def positive_probs_std(self, positive_probs, request):
-        if request.param:
-            std = np.full_like(positive_probs, 0.1)
-            std[positive_probs == 0] = 0
-            std[positive_probs == 1] = 0
-            return std
-        else:
-            return None
-
-    @pytest.fixture(params=[None, "importance", "explained"])
-    def phi_args_base(self, request, input_shape):
-        if input_shape[-1] < 2 and request.param is not None:
-            pytest.skip("not compatible")
-        elif request.param is None:
-            return {}
-        elif request.param == "importance":
-            return {"variables_importance": np.full(input_shape, 2)}
-        else:
-            val = np.full(input_shape, 2)
-            return {"variance_explained": val / val.sum(-1, keepdims=True)}
-
-    @pytest.fixture(params=["concentration", "no-concentration"])
-    def phi_args(self, request, phi_args_base):
-        if request.param == "concentration":
-            phi_args_base["importance_concentration"] = 10
-        return phi_args_base
-
-    def test_init_r2(
-        self,
-        dims,
-        input_std,
-        output_std,
-        r2,
-        r2_std,
-        model: pm.Model,
-    ):
-        eps, beta = pmx.distributions.R2D2M2CP(
-            "beta",
-            output_std,
-            input_std,
-            dims=dims,
-            r2=r2,
-            r2_std=r2_std,
-        )
-        assert not np.isnan(beta.eval()).any()
-        assert eps.eval().shape == output_std.shape
-        assert beta.eval().shape == input_std.shape
-        # r2 rv is only created if r2 std is not None
-        assert "beta" in model.named_vars
-        assert ("beta::r2" in model.named_vars) == (r2_std is not None), set(model.named_vars)
-        # phi is only created if variable importance is not None and there is more than one var
-        assert np.isfinite(model.compile_logp()(model.initial_point()))
-
-    def test_init_importance(
-        self,
-        dims,
-        centered,
-        input_std,
-        output_std,
-        phi_args,
-        model: pm.Model,
-    ):
-        eps, beta = pmx.distributions.R2D2M2CP(
-            "beta",
-            output_std,
-            input_std,
-            dims=dims,
-            r2=1,
-            centered=centered,
-            **phi_args,
-        )
-        assert not np.isnan(beta.eval()).any()
-        assert eps.eval().shape == output_std.shape
-        assert beta.eval().shape == input_std.shape
-        # r2 rv is only created if r2 std is not None
-        assert "beta" in model.named_vars
-        # phi is only created if variable importance is not None and there is more than one var
-        assert ("beta::phi" in model.named_vars) == (
-            "variables_importance" in phi_args or "importance_concentration" in phi_args
-        ), set(model.named_vars)
-        assert np.isfinite(model.compile_logp()(model.initial_point()))
-
-    def test_init_positive_probs(
-        self,
-        dims,
-        centered,
-        input_std,
-        output_std,
-        positive_probs,
-        positive_probs_std,
-        model: pm.Model,
-    ):
-        eps, beta = pmx.distributions.R2D2M2CP(
-            "beta",
-            output_std,
-            input_std,
-            dims=dims,
-            r2=1.0,
-            centered=centered,
-            positive_probs_std=positive_probs_std,
-            positive_probs=positive_probs,
-        )
-        assert not np.isnan(beta.eval()).any()
-        assert eps.eval().shape == output_std.shape
-        assert beta.eval().shape == input_std.shape
-        # r2 rv is only created if r2 std is not None
-        assert "beta" in model.named_vars
-        # phi is only created if variable importance is not None and there is more than one var
-        assert ("beta::psi" in model.named_vars) == (
-            positive_probs_std is not None and positive_probs_std.any()
-        ), set(model.named_vars)
-        assert np.isfinite(sum(model.point_logps().values()))
-
-    def test_failing_importance(self, dims, input_shape, output_std, input_std):
-        if input_shape[-1] < 2:
-            with pytest.raises(TypeError, match="less than two variables"):
-                pmx.distributions.R2D2M2CP(
-                    "beta",
-                    output_std,
-                    input_std,
-                    dims=dims,
-                    r2=0.8,
-                    variables_importance=abs(input_std),
-                )
-        else:
-            pmx.distributions.R2D2M2CP(
-                "beta",
-                output_std,
-                input_std,
-                dims=dims,
-                r2=0.8,
-                variables_importance=abs(input_std),
-            )
-
-    def test_failing_variance_explained(self, dims, input_shape, output_std, input_std):
-        if input_shape[-1] < 2:
-            with pytest.raises(TypeError, match="less than two variables"):
-                pmx.distributions.R2D2M2CP(
-                    "beta",
-                    output_std,
-                    input_std,
-                    dims=dims,
-                    r2=0.8,
-                    variance_explained=abs(input_std),
-                )
-        else:
-            pmx.distributions.R2D2M2CP(
-                "beta", output_std, input_std, dims=dims, r2=0.8, variance_explained=abs(input_std)
-            )
-
-    def test_failing_mutual_exclusive(self, model: pm.Model):
-        with pytest.raises(TypeError, match="variable importance with variance explained"):
-            with model:
-                model.add_coord("a", range(2))
-            pmx.distributions.R2D2M2CP(
-                "beta",
-                1,
-                [1, 1],
-                dims="a",
-                r2=0.8,
-                variance_explained=[0.5, 0.5],
-                variables_importance=[1, 1],
-            )
-
-    def test_limit_case_requires_std_0(self, model: pm.Model):
-        model.add_coord("a", range(2))
-        with pytest.raises(ValueError, match="Can't have both positive_probs"):
-            pmx.distributions.R2D2M2CP(
-                "beta",
-                1,
-                [1, 1],
-                dims="a",
-                r2=0.8,
-                positive_probs=[0.5, 0],
-                positive_probs_std=[0.3, 0.1],
-            )
-        with pytest.raises(ValueError, match="Can't have both positive_probs"):
-            pmx.distributions.R2D2M2CP(
-                "beta",
-                1,
-                [1, 1],
-                dims="a",
-                r2=0.8,
-                positive_probs=[0.5, 1],
-                positive_probs_std=[0.3, 0.1],
-            )
-
-    def test_limit_case_creates_masked_vars(self, model: pm.Model, centered: bool):
-        model.add_coord("a", range(2))
-        pmx.distributions.R2D2M2CP(
-            "beta0",
-            1,
-            [1, 1],
-            dims="a",
-            r2=0.8,
-            positive_probs=[0.5, 1],
-            positive_probs_std=[0.3, 0],
-            centered=centered,
-        )
-        pmx.distributions.R2D2M2CP(
-            "beta1",
-            1,
-            [1, 1],
-            dims="a",
-            r2=0.8,
-            positive_probs=[0.5, 0],
-            positive_probs_std=[0.3, 0],
-            centered=centered,
-        )
-        if not centered:
-            assert "beta0::raw::masked" in model.named_vars, model.named_vars
-            assert "beta1::raw::masked" in model.named_vars, model.named_vars
-        else:
-            assert "beta0::masked" in model.named_vars, model.named_vars
-            assert "beta1::masked" in model.named_vars, model.named_vars
-        assert "beta1::psi::masked" in model.named_vars
-        assert "beta0::psi::masked" in model.named_vars
-
-    def test_zero_length_rvs_not_created(self, model: pm.Model):
-        model.add_coord("a", range(2))
-        # deterministic case which should not have any new variables
-        b = pmx.distributions.R2D2M2CP("b1", 1, [1, 1], r2=0.5, positive_probs=[1, 1], dims="a")
-        assert not model.free_RVs, model.free_RVs
-
-        b = pmx.distributions.R2D2M2CP(
-            "b2", 1, [1, 1], r2=0.5, positive_probs=[1, 1], positive_probs_std=[0, 0], dims="a"
-        )
-        assert not model.free_RVs, model.free_RVs

tests/distributions/test_transform.py

@@ -1,77 +0,0 @@
-#   Copyright 2025 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
-
-from pymc_extras.distributions.transforms import PartialOrder
-
-
-class TestPartialOrder:
-    adj_mats = np.array(
-        [
-            # 0 < {1, 2} < 3
-            [[0, 1, 1, 0], [0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 0]],
-            # 1 < 0 < 3 < 2
-            [[0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0]],
-        ]
-    )
-
-    valid_values = np.array([[0, 2, 1, 3], [1, 0, 3, 2]], dtype=float)
-
-    # Test that forward and backward are inverses of eachother
-    # And that it works when extra dimensions are added in data
-    def test_forward_backward_dimensionality(self):
-        po = PartialOrder(self.adj_mats)
-        po0 = PartialOrder(self.adj_mats[0])
-        vv = self.valid_values
-        vv0 = self.valid_values[0]
-
-        testsets = [
-            (vv, po),
-            (po.initvals(), po),
-            (vv0, po0),
-            (po0.initvals(), po0),
-            (np.tile(vv0, (2, 1)), po0),
-            (np.tile(vv0, (2, 3, 2, 1)), po0),
-            (np.tile(vv, (2, 3, 2, 1, 1)), po),
-        ]
-
-        for vv, po in testsets:
-            fw = po.forward(vv)
-            bw = po.backward(fw)
-            np.testing.assert_allclose(bw.eval(), vv)
-
-    def test_sample_model(self):
-        po = PartialOrder(self.adj_mats)
-        with pm.Model() as model:
-            x = pm.Normal(
-                "x",
-                size=(3, 2, 4),
-                transform=po,
-                initval=po.initvals(shape=(3, 2, 4), lower=-1, upper=1),
-            )
-            idata = pm.sample()
-
-        # Check that the order constraints are satisfied
-        # Move chain, draw and "3" dimensions to the back
-        xvs = idata.posterior.x.values.transpose(3, 4, 0, 1, 2)
-        x0 = xvs[0]  # 0 < {1, 2} < 3
-        assert (
-            (x0[0] < x0[1]).all()
-            and (x0[0] < x0[2]).all()
-            and (x0[1] < x0[3]).all()
-            and (x0[2] < x0[3]).all()
-        )
-        x1 = xvs[1]  # 1 < 0 < 3 < 2
-        assert (x1[1] < x1[0]).all() and (x1[0] < x1[3]).all() and (x1[3] < x1[2]).all()

tests/model/marginal/test_distributions.py

@@ -1,132 +0,0 @@
-import numpy as np
-import pymc as pm
-import pytest
-
-from pymc.logprob.abstract import _logprob
-from pytensor import tensor as pt
-from scipy.stats import norm
-
-from pymc_extras import marginalize
-from pymc_extras.distributions import DiscreteMarkovChain
-from pymc_extras.model.marginal.distributions import MarginalFiniteDiscreteRV
-
-
-def test_marginalized_bernoulli_logp():
-    """Test logp of IR TestFiniteMarginalDiscreteRV directly"""
-    mu = pt.vector("mu")
-
-    idx = pm.Bernoulli.dist(0.7, name="idx")
-    y = pm.Normal.dist(mu=mu[idx], sigma=1.0, name="y")
-    marginal_rv_node = MarginalFiniteDiscreteRV(
-        [mu],
-        [idx, y],
-        dims_connections=(((),),),
-        dims=(),
-    )(mu)[0].owner
-
-    y_vv = y.clone()
-    (logp,) = _logprob(
-        marginal_rv_node.op,
-        (y_vv,),
-        *marginal_rv_node.inputs,
-    )
-
-    ref_logp = pm.logp(pm.NormalMixture.dist(w=[0.3, 0.7], mu=mu, sigma=1.0), y_vv)
-    np.testing.assert_almost_equal(
-        logp.eval({mu: [-1, 1], y_vv: 2}),
-        ref_logp.eval({mu: [-1, 1], y_vv: 2}),
-    )
-
-
-@pytest.mark.parametrize("batch_chain", (False, True), ids=lambda x: f"batch_chain={x}")
-@pytest.mark.parametrize("batch_emission", (False, True), ids=lambda x: f"batch_emission={x}")
-def test_marginalized_hmm_normal_emission(batch_chain, batch_emission):
-    if batch_chain and not batch_emission:
-        pytest.skip("Redundant implicit combination")
-
-    with pm.Model() as m:
-        P = [[0, 1], [1, 0]]
-        init_dist = pm.Categorical.dist(p=[1, 0])
-        chain = DiscreteMarkovChain(
-            "chain", P=P, init_dist=init_dist, steps=3, shape=(3, 4) if batch_chain else None
-        )
-        emission = pm.Normal(
-            "emission", mu=chain * 2 - 1, sigma=1e-1, shape=(3, 4) if batch_emission else None
-        )
-
-    marginal_m = marginalize(m, [chain])
-    logp_fn = marginal_m.compile_logp()
-
-    test_value = np.array([-1, 1, -1, 1])
-    expected_logp = pm.logp(pm.Normal.dist(0, 1e-1), np.zeros_like(test_value)).sum().eval()
-    if batch_emission:
-        test_value = np.broadcast_to(test_value, (3, 4))
-        expected_logp *= 3
-    np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp)
-
-
-@pytest.mark.parametrize(
-    "categorical_emission",
-    [False, True],
-)
-def test_marginalized_hmm_categorical_emission(categorical_emission):
-    """Example adapted from https://www.youtube.com/watch?v=9-sPm4CfcD0"""
-    with pm.Model() as m:
-        P = np.array([[0.5, 0.5], [0.3, 0.7]])
-        init_dist = pm.Categorical.dist(p=[0.375, 0.625])
-        chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, steps=2)
-        if categorical_emission:
-            emission = pm.Categorical("emission", p=pt.constant([[0.8, 0.2], [0.4, 0.6]])[chain])
-        else:
-            emission = pm.Bernoulli("emission", p=pt.where(pt.eq(chain, 0), 0.2, 0.6))
-    marginal_m = marginalize(m, [chain])
-
-    test_value = np.array([0, 0, 1])
-    expected_logp = np.log(0.1344)  # Shown at the 10m22s mark in the video
-    logp_fn = marginal_m.compile_logp()
-    np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp)
-
-
-@pytest.mark.parametrize("batch_chain", (False, True))
-@pytest.mark.parametrize("batch_emission1", (False, True))
-@pytest.mark.parametrize("batch_emission2", (False, True))
-def test_marginalized_hmm_multiple_emissions(batch_chain, batch_emission1, batch_emission2):
-    chain_shape = (3, 1, 4) if batch_chain else (4,)
-    emission1_shape = (
-        (2, *reversed(chain_shape)) if batch_emission1 else tuple(reversed(chain_shape))
-    )
-    emission2_shape = (*chain_shape, 2) if batch_emission2 else chain_shape
-    with pm.Model() as m:
-        P = [[0, 1], [1, 0]]
-        init_dist = pm.Categorical.dist(p=[1, 0])
-        chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, shape=chain_shape)
-        emission_1 = pm.Normal(
-            "emission_1", mu=(chain * 2 - 1).T, sigma=1e-1, shape=emission1_shape
-        )
-
-        emission2_mu = (1 - chain) * 2 - 1
-        if batch_emission2:
-            emission2_mu = emission2_mu[..., None]
-        emission_2 = pm.Normal("emission_2", mu=emission2_mu, sigma=1e-1, shape=emission2_shape)
-
-    marginal_m = marginalize(m, [chain])
-
-    with pytest.warns(UserWarning, match="multiple dependent variables"):
-        logp_fn = marginal_m.compile_logp(sum=False)
-
-    test_value = np.array([-1, 1, -1, 1])
-    multiplier = 2 + batch_emission1 + batch_emission2
-    if batch_chain:
-        multiplier *= 3
-    expected_logp = norm.logpdf(np.zeros_like(test_value), 0, 1e-1).sum() * multiplier
-
-    test_value = np.broadcast_to(test_value, chain_shape)
-    test_value_emission1 = np.broadcast_to(test_value.T, emission1_shape)
-    if batch_emission2:
-        test_value_emission2 = np.broadcast_to(-test_value[..., None], emission2_shape)
-    else:
-        test_value_emission2 = np.broadcast_to(-test_value, emission2_shape)
-    test_point = {"emission_1": test_value_emission1, "emission_2": test_value_emission2}
-    res_logp, dummy_logp = logp_fn(test_point)
-    assert res_logp.shape == ((1, 3) if batch_chain else ())
-    np.testing.assert_allclose(res_logp.sum(), expected_logp)

tests/model/marginal/test_graph_analysis.py

@@ -1,182 +0,0 @@
-import pytensor.tensor as pt
-import pytest
-
-from pymc.distributions import CustomDist
-from pytensor.tensor.type_other import NoneTypeT
-
-from pymc_extras.model.marginal.graph_analysis import (
-    is_conditional_dependent,
-    subgraph_batch_dim_connection,
-)
-
-
-def test_is_conditional_dependent_static_shape():
-    """Test that we don't consider dependencies through "constant" shape Ops"""
-    x1 = pt.matrix("x1", shape=(None, 5))
-    y1 = pt.random.normal(size=pt.shape(x1))
-    assert is_conditional_dependent(y1, x1, [x1, y1])
-
-    x2 = pt.matrix("x2", shape=(9, 5))
-    y2 = pt.random.normal(size=pt.shape(x2))
-    assert not is_conditional_dependent(y2, x2, [x2, y2])
-
-
-class TestSubgraphBatchDimConnection:
-    def test_dimshuffle(self):
-        inp = pt.tensor(shape=(5, 1, 4, 3))
-        out1 = pt.matrix_transpose(inp)
-        out2 = pt.expand_dims(inp, 1)
-        out3 = pt.squeeze(inp)
-        [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [out1, out2, out3])
-        assert dims1 == (0, 1, 3, 2)
-        assert dims2 == (0, None, 1, 2, 3)
-        assert dims3 == (0, 2, 3)
-
-    def test_careduce(self):
-        inp = pt.tensor(shape=(4, 3, 2))
-
-        out = pt.sum(inp[:, None], axis=(1,))
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, 2)
-
-        invalid_out = pt.sum(inp, axis=(1,))
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-    def test_subtensor(self):
-        inp = pt.tensor(shape=(4, 3, 2))
-
-        invalid_out = inp[0, :1]
-        with pytest.raises(
-            ValueError,
-            match="Partial slicing or indexing of known dimensions not supported",
-        ):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-        # If we are selecting dummy / unknown dimensions that's fine
-        valid_out = pt.expand_dims(inp, (0, 1))[0, :1]
-        [dims] = subgraph_batch_dim_connection(inp, [valid_out])
-        assert dims == (None, 0, 1, 2)
-
-    def test_advanced_subtensor_value(self):
-        inp = pt.tensor(shape=(2, 4))
-        intermediate_out = inp[:, None, :, None] + pt.zeros((2, 3, 4, 5))
-
-        # Index on an unlabled dim introduced by broadcasting with zeros
-        out = intermediate_out[:, [0, 0, 1, 2]]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, None, 1, None)
-
-        # Indexing that introduces more dimensions
-        out = intermediate_out[:, [[0, 0], [1, 2]], :]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, None, None, 1, None)
-
-        # Special case where advanced dims are moved to the front of the output
-        out = intermediate_out[:, [0, 0, 1, 2], :, 0]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (None, 0, 1)
-
-        # Indexing on a labeled dim fails
-        out = intermediate_out[:, :, [0, 0, 1, 2]]
-        with pytest.raises(ValueError, match="Partial slicing or advanced integer indexing"):
-            subgraph_batch_dim_connection(inp, [out])
-
-    def test_advanced_subtensor_key(self):
-        inp = pt.tensor(shape=(5, 5), dtype=int)
-        base = pt.zeros((2, 3, 4))
-
-        out = base[inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, None, None)
-
-        out = base[:, :, inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (
-            None,
-            None,
-            0,
-            1,
-        )
-
-        out = base[1:, 0, inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (None, 0, 1)
-
-        # Special case where advanced dims are moved to the front of the output
-        out = base[0, :, inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, None)
-
-        # Mix keys dimensions
-        out = base[:, inp, inp.T]
-        with pytest.raises(ValueError, match="Different known dimensions mixed via broadcasting"):
-            subgraph_batch_dim_connection(inp, [out])
-
-    def test_elemwise(self):
-        inp = pt.tensor(shape=(5, 5))
-
-        out = inp + inp
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1)
-
-        out = inp + inp.T
-        with pytest.raises(ValueError, match="Different known dimensions mixed via broadcasting"):
-            subgraph_batch_dim_connection(inp, [out])
-
-        out = inp[None, :, None, :] + inp[:, None, :, None]
-        with pytest.raises(
-            ValueError, match="Same known dimension used in different axis after broadcasting"
-        ):
-            subgraph_batch_dim_connection(inp, [out])
-
-    def test_blockwise(self):
-        inp = pt.tensor(shape=(5, 4))
-
-        invalid_out = inp @ pt.ones((4, 3))
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-        out = (inp[:, :, None, None] + pt.zeros((2, 3))) @ pt.ones((2, 3))
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, None, None)
-
-    def test_random_variable(self):
-        inp = pt.tensor(shape=(5, 4, 3))
-
-        out1 = pt.random.normal(loc=inp)
-        out2 = pt.random.categorical(p=inp[..., None])
-        out3 = pt.random.multivariate_normal(mean=inp[..., None], cov=pt.eye(1))
-        [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [out1, out2, out3])
-        assert dims1 == (0, 1, 2)
-        assert dims2 == (0, 1, 2)
-        assert dims3 == (0, 1, 2, None)
-
-        invalid_out = pt.random.categorical(p=inp)
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-        invalid_out = pt.random.multivariate_normal(mean=inp, cov=pt.eye(3))
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-    def test_symbolic_random_variable(self):
-        inp = pt.tensor(shape=(4, 3, 2))
-
-        # Test univariate
-        out = CustomDist.dist(
-            inp,
-            dist=lambda mu, size: pt.random.normal(loc=mu, size=size),
-        )
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, 2)
-
-        # Test multivariate
-        def dist(mu, size):
-            if isinstance(size.type, NoneTypeT):
-                size = mu.shape
-            return pt.random.normal(loc=mu[..., None], size=(*size, 2))
-
-        out = CustomDist.dist(inp, dist=dist, size=(4, 3, 2), signature="()->(2)")
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, 2, None)