Add test for univariate and multivariate marginal mixture

Co-authored-by: Jesse Grabowski <[email protected]>

Author: ricardoV94

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))