.wip test

Author: ricardoV94

pymc_experimental/model/marginal/marginal_model.py

@@ -592,7 +592,6 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
             extra_batch_ndim = dependent_rv.type.ndim + dependent_rv_ndim_supp - marginal_ndim
             valid_dependent_batch_dims = marginal_batch_dims + (((),) * extra_batch_ndim)
             if dependent_rv_batch_dims != valid_dependent_batch_dims:
-                # TODO: This message is too specific
                 raise NotImplementedError(
                     f"Link between dimensions of marginalized and dependent RVs not supported: {dependent_rv_batch_dims} != {valid_dependent_batch_dims}"
                 )

tests/model/marginal/test_graph_analysis.py

@@ -4,7 +4,21 @@
 from pymc.distributions import CustomDist
 from pytensor.tensor.type_other import NoneTypeT
 
-from pymc_experimental.model.marginal.graph_analysis import subgraph_batch_dim_connection
+from pymc_experimental.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:

tests/model/marginal/test_marginal_model.py

@@ -17,15 +17,14 @@
 from scipy.special import log_softmax, logsumexp
 from scipy.stats import halfnorm, norm
 
-from pymc_experimental.model.marginal.graph_analysis import is_conditional_dependent
 from pymc_experimental.model.marginal.marginal_model import (
     MarginalModel,
     marginalize,
 )
 from tests.utils import equal_computations_up_to_root
 
 
-def test_marginalized_basic():
+def test_basic_marginalized_rv():
     data = [2] * 5
 
     with MarginalModel() as m:
@@ -69,7 +68,8 @@ def test_marginalized_basic():
         )
 
 
-def test_multiple_independent_marginalized_rvs():
+def test_one_to_one_marginalized_rvs():
+    """Test case with multiple, independent marginalized RVs"""
     with MarginalModel() as m:
         sigma = pm.HalfNormal("sigma")
         idx1 = pm.Bernoulli("idx1", p=0.75)
@@ -95,7 +95,7 @@ def test_multiple_independent_marginalized_rvs():
     np.testing.assert_array_almost_equal(y_logp, y_ref_logp)
 
 
-def test_multiple_dependent_marginalized_rvs():
+def test_one_to_many_marginalized_rvs():
     """Test that marginalization works when there is more than one dependent RV"""
     with MarginalModel() as m:
         sigma = pm.HalfNormal("sigma")
@@ -118,7 +118,37 @@ def test_multiple_dependent_marginalized_rvs():
     np.testing.assert_array_almost_equal(logp_x_y, ref_logp_x_y)
 
 
-def test_rv_dependent_multiple_marginalized_rvs():
+def test_one_to_many_unaligned_marginalized_rvs():
+    """Test that marginalization works when there is more than one dependent RV with batch dimensions that are not aligned"""
+
+    def build_model(build_batched: bool):
+        with MarginalModel() as m:
+            if build_batched:
+                idx = pm.Bernoulli("idx", p=[0.75, 0.4], shape=(3, 2))
+            else:
+                idxs = [pm.Bernoulli(f"idx_{i}", p=(0.75 if i % 2 == 0 else 0.4)) for i in range(6)]
+                idx = pt.stack(idxs, axis=0).reshape(3, 2)
+
+            x = pm.Normal("x", mu=idx.T[:, :, None], shape=(2, 3, 1))
+            y = pm.Normal("y", mu=(idx * 2 - 1), shape=(1, 3, 2))
+
+        return m
+
+    m = build_model(build_batched=True)
+    with pytest.warns(UserWarning, match="There are multiple dependent variables"):
+        m.marginalize(["idx"])
+
+    ref_m = build_model(build_batched=False)
+    ref_m.marginalize([f"idx_{i}" for i in range(6)])
+
+    test_point = m.initial_point()
+    np.testing.assert_allclose(
+        m.compile_logp()(test_point),
+        ref_m.compile_logp()(test_point),
+    )
+
+
+def test_many_to_one_marginalized_rvs():
     """Test when random variables depend on multiple marginalized variables"""
     with MarginalModel() as m:
         x = pm.Bernoulli("x", 0.1)
@@ -133,13 +163,13 @@ def test_rv_dependent_multiple_marginalized_rvs():
     np.testing.assert_allclose(np.exp(logp({"z": 2})), 0.1 * 0.3)
 
 
-@pytest.mark.parametrize("batched", (False, True))
+@pytest.mark.parametrize("batched", (False, "left", "right"))
 def test_nested_marginalized_rvs(batched):
     """Test that marginalization works when there are nested marginalized RVs"""
 
     def build_model(build_batched: bool) -> MarginalModel:
         idx_shape = (3,) if build_batched else ()
-        sub_idx_shape = (3, 5) if build_batched else (5,)
+        sub_idx_shape = (5,) if not build_batched else (5, 3) if batched == "left" else (3, 5)
 
         with MarginalModel() as m:
             sigma = pm.HalfNormal("sigma")
@@ -148,9 +178,9 @@ def build_model(build_batched: bool) -> MarginalModel:
             dep = pm.Normal("dep", mu=pt.switch(pt.eq(idx, 0), -1000.0, 1000.0), sigma=sigma)
 
             sub_idx_p = pt.switch(pt.eq(idx, 0), 0.15, 0.95)
-            if build_batched:
-                sub_idx_p = sub_idx_p[:, None]
-                dep = dep[:, None]
+            if build_batched and batched == "right":
+                sub_idx_p = sub_idx_p[..., None]
+                dep = dep[..., None]
             sub_idx = pm.Bernoulli("sub_idx", p=sub_idx_p, shape=sub_idx_shape)
             sub_dep = pm.Normal("sub_dep", mu=dep + sub_idx * 100, sigma=sigma)
 
@@ -204,22 +234,22 @@ def test_marginalized_index_as_key(advanced_indexing):
 
     with MarginalModel() as m:
         x = pm.Categorical("x", p=w, shape=shape)
-        y = pm.Normal("y", mu[x], sigma=1, observed=y_val)
+        y = pm.Normal("y", mu[x].T, sigma=1, observed=y_val)
 
     m.marginalize(x)
 
     marginal_logp = m.compile_logp(sum=False)({})[0]
-    ref_logp = pm.logp(pm.NormalMixture.dist(w=w, mu=mu, sigma=1, shape=shape), y_val).eval()
+    ref_logp = pm.logp(pm.NormalMixture.dist(w=w, mu=mu.T, sigma=1, shape=shape), y_val).eval()
 
     np.testing.assert_allclose(marginal_logp, ref_logp)
 
 
 def test_marginalized_index_as_value_and_key():
     """Test we can marginalize graphs were marginalized_rv is indexed."""
 
-    def build_model(batch: bool) -> MarginalModel:
+    def build_model(build_batched: bool) -> MarginalModel:
         with MarginalModel() as m:
-            if batch:
+            if build_batched:
                 latent_state = pm.Bernoulli("latent_state", p=0.3, size=(4,))
             else:
                 latent_state = pm.math.stack(
@@ -237,8 +267,8 @@ def build_model(batch: bool) -> MarginalModel:
         return m
 
     # We compare with the equivalent but less efficient batched model
-    m = build_model(batch=True)
-    ref_m = build_model(batch=False)
+    m = build_model(build_batched=True)
+    ref_m = build_model(build_batched=False)
 
     m.marginalize(["latent_state"])
     ref_m.marginalize([f"latent_state_{i}" for i in range(4)])
@@ -317,6 +347,14 @@ def test_mixed_dims_via_support_dimension(self):
             with pytest.raises(NotImplementedError):
                 m.marginalize(x)
 
+    def test_mixed_dims_via_nested_marginalization(self):
+        with MarginalModel() as m:
+            x = pm.Bernoulli("x", p=0.7, shape=(3,))
+            y = pm.Bernoulli("y", p=0.7, shape=(2,))
+            z = pm.Normal("z", mu=pt.add.outer(x, y), shape=(3, 2))
+            with pytest.raises(NotImplementedError):
+                m.marginalize([x, y])
+
 
 def test_marginalized_deterministic_and_potential():
     rng = np.random.default_rng(299)
@@ -432,17 +470,6 @@ def test_marginalized_transforms(transform, expected_warning):
     np.testing.assert_allclose(m.compile_logp()(ip), m_ref.compile_logp()(ip))
 
 
-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])
-
-
 def test_data_container():
     """Test that MarginalModel can handle Data containers."""
     with MarginalModel(coords={"obs": [0]}) as marginal_m:
@@ -469,49 +496,6 @@ def test_data_container():
         np.testing.assert_allclose(logp_fn(ip), ref_logp_fn(ip))
 
 
-@pytest.mark.parametrize("univariate", (True, False))
-def test_vector_univariate_mixture(univariate):
-    with MarginalModel() as m:
-        idx = pm.Bernoulli("idx", p=0.5, shape=(2,) if univariate else ())
-
-        def dist(idx, size):
-            return pm.math.switch(
-                pm.math.eq(idx, 0),
-                pm.Normal.dist([-10, -10], 1),
-                pm.Normal.dist([10, 10], 1),
-            )
-
-        pm.CustomDist("norm", idx, dist=dist)
-
-    m.marginalize(idx)
-    logp_fn = m.compile_logp()
-
-    if univariate:
-        with pm.Model() as ref_m:
-            pm.NormalMixture("norm", w=[0.5, 0.5], mu=[[-10, 10], [-10, 10]], shape=(2,))
-    else:
-        with pm.Model() as ref_m:
-            pm.Mixture(
-                "norm",
-                w=[0.5, 0.5],
-                comp_dists=[
-                    pm.MvNormal.dist([-10, -10], np.eye(2)),
-                    pm.MvNormal.dist([10, 10], np.eye(2)),
-                ],
-                shape=(2,),
-            )
-    ref_logp_fn = ref_m.compile_logp()
-
-    for test_value in (
-        [-10, -10],
-        [10, 10],
-        [-10, 10],
-        [-10, 10],
-    ):
-        pt = {"norm": test_value}
-        np.testing.assert_allclose(logp_fn(pt), ref_logp_fn(pt))
-
-
 def test_mutable_indexing_jax_backend():
     pytest.importorskip("jax")
     from pymc.sampling.jax import get_jaxified_logp
@@ -631,11 +615,51 @@ def test_change_point_model_sampling(self, disaster_model):
             rtol=1e-2,
         )
 
-    @pytest.mark.parametrize(
-        "batch_right", (True, pytest.param(False, marks=pytest.mark.xfail(reason="NotImplemented")))
-    )
+    @pytest.mark.parametrize("univariate", (True, False))
+    def test_vector_univariate_mixture(self, univariate):
+        with MarginalModel() as m:
+            idx = pm.Bernoulli("idx", p=0.5, shape=(2,) if univariate else ())
+
+            def dist(idx, size):
+                return pm.math.switch(
+                    pm.math.eq(idx, 0),
+                    pm.Normal.dist([-10, -10], 1),
+                    pm.Normal.dist([10, 10], 1),
+                )
+
+            pm.CustomDist("norm", idx, dist=dist)
+
+        m.marginalize(idx)
+        logp_fn = m.compile_logp()
+
+        if univariate:
+            with pm.Model() as ref_m:
+                pm.NormalMixture("norm", w=[0.5, 0.5], mu=[[-10, 10], [-10, 10]], shape=(2,))
+        else:
+            with pm.Model() as ref_m:
+                pm.Mixture(
+                    "norm",
+                    w=[0.5, 0.5],
+                    comp_dists=[
+                        pm.MvNormal.dist([-10, -10], np.eye(2)),
+                        pm.MvNormal.dist([10, 10], np.eye(2)),
+                    ],
+                    shape=(2,),
+                )
+        ref_logp_fn = ref_m.compile_logp()
+
+        for test_value in (
+            [-10, -10],
+            [10, 10],
+            [-10, 10],
+            [-10, 10],
+        ):
+            pt = {"norm": test_value}
+            np.testing.assert_allclose(logp_fn(pt), ref_logp_fn(pt))
+
+    @pytest.mark.parametrize("batch_right", (True, False))
     def test_k_censored_clusters_model(self, batch_right):
-        def build_model(batch: bool) -> MarginalModel:
+        def build_model(build_batched: bool) -> MarginalModel:
             data = np.array([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])
             nobs = data.shape[0]
             n_clusters = 5
@@ -645,7 +669,7 @@ def build_model(batch: bool) -> MarginalModel:
                 "obs": range(nobs),
             }
             with MarginalModel(coords=coords) as m:
-                if batch:
+                if build_batched:
                     idx = pm.Categorical("idx", p=np.ones(n_clusters) / n_clusters, dims=["obs"])
                 else:
                     idx = pm.math.stack(
@@ -682,8 +706,8 @@ def build_model(batch: bool) -> MarginalModel:
 
             return m
 
-        m = build_model(batch=True)
-        ref_m = build_model(batch=False)
+        m = build_model(build_batched=True)
+        ref_m = build_model(build_batched=False)
 
         m.marginalize([m["idx"]])
         ref_m.marginalize([n for n in ref_m.named_vars if n.startswith("idx_")])