Restore ndim_supp check, add assumed ndim_supp based on the type of RV being marginalized

Author: jessegrabowski

pymc_experimental/marginal_model.py

@@ -349,11 +349,9 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     ndim_supp = {rv.owner.op.ndim_supp for rv in dependent_rvs}
     if len(ndim_supp) != 1:
         raise NotImplementedError()
-    ndim_supp = tuple(ndim_supp)[0]
-    # if max(ndim_supp) > 0:
-    #     raise NotImplementedError(
-    #         "Marginalization with dependent Multivariate RVs not implemented"
-    #     )
+
+    if max(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 = [
@@ -390,13 +388,15 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
 
     if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
         marginalize_constructor = DiscreteMarginalMarkovChainRV
+        marginalize_target_rv_ndim_supp = 0
     else:
         marginalize_constructor = FiniteDiscreteMarginalRV
+        marginalize_target_rv_ndim_supp = 1
 
     marginalization_op = marginalize_constructor(
         inputs=list(replace_inputs.values()),
         outputs=cloned_outputs,
-        ndim_supp=ndim_supp,
+        ndim_supp=marginalize_target_rv_ndim_supp,
     )
 
     marginalized_rvs = marginalization_op(*replace_inputs.keys())