Add test for batched marginalized variables

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -368,14 +368,14 @@ def transform_input(inputs):
 
             rv_shape = constant_fold(tuple(rv.shape))
             rv_domain = get_domain_of_finite_discrete_rv(rv)
-            rv_domain_tensor = pt.swapaxes(
+            rv_domain_tensor = pt.moveaxis(
                 pt.full(
                     (*rv_shape, len(rv_domain)),
                     rv_domain,
                     dtype=rv.dtype,
                 ),
-                axis1=0,
-                axis2=-1,
+                -1,
+                0,
             )
 
             marginalized_value = m.rvs_to_values[rv]
@@ -641,14 +641,14 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
     # PyMC does not allow RVs in the logp graph, even if we are just using the shape
     marginalized_rv_shape = constant_fold(tuple(marginalized_rv.shape))
     marginalized_rv_domain = get_domain_of_finite_discrete_rv(marginalized_rv)
-    marginalized_rv_domain_tensor = pt.swapaxes(
+    marginalized_rv_domain_tensor = pt.moveaxis(
         pt.full(
             (*marginalized_rv_shape, len(marginalized_rv_domain)),
             marginalized_rv_domain,
             dtype=marginalized_rv.dtype,
         ),
-        axis1=0,
-        axis2=-1,
+        -1,
+        0,
     )
 
     # Arbitrary cutoff to switch to Scan implementation to keep graph size under control

pymc_experimental/tests/model/test_marginal_model.py

@@ -297,6 +297,30 @@ def true_logp(y, sigma):
     )
 
 
+def test_recover_batched_marginal():
+    """Test that marginalization works for batched random variables"""
+    with MarginalModel() as m:
+        sigma = pm.HalfNormal("sigma")
+        idx = pm.Bernoulli("idx", p=0.7, shape=(2, 2))
+        y = pm.Normal("y", mu=idx, sigma=sigma, shape=(2, 2))
+
+    m.marginalize([idx])
+
+    rng = np.random.default_rng(211)
+
+    with m:
+        prior = pm.sample_prior_predictive(
+            samples=20,
+            random_seed=rng,
+            return_inferencedata=False,
+        )
+        idata = InferenceData(
+            posterior=dict_to_dataset({k: np.expand_dims(prior[k], axis=0) for k in prior})
+        )
+
+    idata = m.recover_marginals(idata, return_samples=True)
+
+
 def test_nested_recover_marginals():
     """Test that marginalization works when there are nested marginalized RVs"""