.WIP refactor MarginalModel

pymc_experimental/model/marginal/distributions.py

@@ -4,17 +4,19 @@
 import pytensor.tensor as pt
 
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
+from pymc.distributions.distribution import _support_point, support_point
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
 from pytensor import Variable
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.mode import Mode
-from pytensor.graph import Op, vectorize_graph
+from pytensor.graph import FunctionGraph, Op, vectorize_graph
 from pytensor.graph.replace import clone_replace, graph_replace
 from pytensor.scan import map as scan_map
 from pytensor.scan import scan
 from pytensor.tensor import TensorVariable
+from pytensor.tensor.random.type import RandomType
 
 from pymc_experimental.distributions import DiscreteMarkovChain
 
@@ -44,6 +46,56 @@ def support_axes(self) -> tuple[tuple[int]]:
         return tuple(support_axes_vars)
 
 
+@_support_point.register
+def support_point_marginal_rv(op: MarginalRV, rv, *inputs):
+    """Support point for a marginalized RV.
+
+    The support point of a marginalized RV is the support point of the inner RV,
+    conditioned on the marginalized RV and other dependent RVs taking their support point.
+    """
+    outputs = rv.owner.outputs
+
+    inner_rv = op.inner_outputs[outputs.index(rv)]
+    marginalized_inner_rv, *other_dependent_inner_rvs = (
+        out
+        for out in op.inner_outputs
+        if out is not inner_rv and not isinstance(out.type, RandomType)
+    )
+
+    # Replace references to inner rvs by the dummy variables (including the marginalized RV)
+    # This is necessary because the inner RVs may depend on each other
+    marginalized_inner_rv_dummy = marginalized_inner_rv.clone()
+    other_dependent_inner_rv_to_dummies = {
+        inner_rv: inner_rv.clone() for inner_rv in other_dependent_inner_rvs
+    }
+    inner_rv = clone_replace(
+        inner_rv,
+        replace={marginalized_inner_rv: marginalized_inner_rv_dummy}
+        | other_dependent_inner_rv_to_dummies,
+    )
+
+    # Get support point of inner RV and marginalized RV
+    inner_rv_support_point = support_point(inner_rv)
+    marginalized_inner_rv_support_point = support_point(marginalized_inner_rv)
+
+    replacements = [
+        # Replace the marginalized RV dummy by its support point
+        (marginalized_inner_rv_dummy, marginalized_inner_rv_support_point),
+        # Replace other dependent RVs dummies by the respective outer outputs.
+        # PyMC will replace them by their support points later
+        *(
+            (v, outputs[op.inner_outputs.index(k)])
+            for k, v in other_dependent_inner_rv_to_dummies.items()
+        ),
+        # Replace outer input RVs
+        *zip(op.inner_inputs, inputs),
+    ]
+    fgraph = FunctionGraph(outputs=[inner_rv_support_point], clone=False)
+    fgraph.replace_all(replacements, import_missing=True)
+    [rv_support_point] = fgraph.outputs
+    return rv_support_point
+
+
 class MarginalFiniteDiscreteRV(MarginalRV):
     """Base class for Marginalized Finite Discrete RVs"""
 

pymc_experimental/model/marginal/graph_analysis.py

@@ -4,6 +4,7 @@
 from itertools import zip_longest
 
 from pymc import SymbolicRandomVariable
+from pymc.model.fgraph import ModelVar
 from pytensor.compile import SharedVariable
 from pytensor.graph import Constant, Variable, ancestors
 from pytensor.graph.basic import io_toposort
@@ -35,12 +36,12 @@ def static_shape_ancestors(vars):
 
 def find_conditional_input_rvs(output_rvs, all_rvs):
     """Find conditionally indepedent input RVs."""
-    blockers = [other_rv for other_rv in all_rvs if other_rv not in output_rvs]
-    blockers += static_shape_ancestors(tuple(all_rvs) + tuple(output_rvs))
+    other_rvs = [other_rv for other_rv in all_rvs if other_rv not in output_rvs]
+    blockers = other_rvs + static_shape_ancestors(tuple(all_rvs) + tuple(output_rvs))
     return [
         var
         for var in ancestors(output_rvs, blockers=blockers)
-        if var in blockers or (var.owner is None and not isinstance(var, Constant | SharedVariable))
+        if var in other_rvs
     ]
 
 
@@ -141,6 +142,9 @@ def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars)
             # None of the inputs are related to the batch_axes of the input_vars
             continue
 
+        elif isinstance(node.op, ModelVar):
+            var_dims[node.outputs[0]] = inputs_dims[0]
+
         elif isinstance(node.op, DimShuffle):
             [input_dims] = inputs_dims
             output_dims = tuple(None if i == "x" else input_dims[i] for i in node.op.new_order)