Minor cleanup MarginalModel

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

@@ -467,7 +467,7 @@ def test_not_supported_marginalized():
         y = pm.Dirichlet("y", a=pm.math.switch(x, [1, 1, 1], [10, 10, 10]))
         with pytest.raises(
             NotImplementedError,
-            match="Marginalization of withe dependent Multivariate RVs not implemented",
+            match="Marginalization with dependent Multivariate RVs not implemented",
         ):
             m.marginalize(x)