Add test for univariate and multivariate marginal mixture

Fix issue with `ndim_supp`

Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

Author: ricardoV94

pymc_experimental/model/marginal_model.py

@@ -580,10 +580,13 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
         raise ValueError(f"No RVs depend on marginalized RV {rv_to_marginalize}")
 
     ndim_supp = {rv.owner.op.ndim_supp for rv in dependent_rvs}
-    if max(ndim_supp) > 0:
+    if len(ndim_supp) != 1:
         raise NotImplementedError(
-            "Marginalization of withe dependent Multivariate RVs not implemented"
+            "Marginalization with dependent variables of different support not implemented"
         )
+    [ndim_supp] = ndim_supp
+    if ndim_supp > 0:
+        raise NotImplementedError("Marginalization with dependent Multivariate RVs not implemented")
 
     marginalized_rv_input_rvs = find_conditional_input_rvs([rv_to_marginalize], all_rvs)
     dependent_rvs_input_rvs = [
@@ -621,7 +624,7 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     marginalization_op = FiniteDiscreteMarginalRV(
         inputs=list(replace_inputs.values()),
         outputs=cloned_outputs,
-        ndim_supp=0,
+        ndim_supp=ndim_supp,
     )
     marginalized_rvs = marginalization_op(*replace_inputs.keys())
     fgraph.replace_all(tuple(zip(rvs_to_marginalize, marginalized_rvs)))

pymc_experimental/tests/model/test_marginal_model.py

@@ -598,3 +598,47 @@ def test_is_conditional_dependent_static_shape():
     x2 = pt.matrix("x2", shape=(9, 5))
     y2 = pt.random.normal(size=pt.shape(x2))
     assert not is_conditional_dependent(y2, x2, [x2, y2])
+
+
+@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))