.WIP refactor MarginalModel

pymc_experimental/model/marginal/distributions.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytensor.tensor as pt
 
+from pymc.distributions.distribution import _support_point
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
@@ -44,6 +45,17 @@ def support_axes(self) -> tuple[tuple[int]]:
         return tuple(support_axes_vars)
 
 
+@_support_point.register
+def support_point_marginal_rv(op: MarginalRV, rv, *inputs):
+    rv_idx = rv.owner.outputs.index(rv)
+    inner_rv = inline_ofg_outputs(op, inputs)[rv_idx]
+    return _support_point(inner_rv.owner.op, inner_rv, *inner_rv.owner.inputs)
+
+
+
+
+
+
 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)