.WIP

Author: ricardoV94

pymc_experimental/model/marginal/distributions.py

@@ -4,11 +4,12 @@
 import pytensor.tensor as pt
 
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform, SymbolicRandomVariable
-from pymc.logprob.abstract import _logprob
+from pymc.logprob.abstract import _logprob, MeasurableOp
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
+from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.mode import Mode
-from pytensor.graph import vectorize_graph
+from pytensor.graph import vectorize_graph, Op
 from pytensor.graph.replace import clone_replace, graph_replace
 from pytensor.scan import map as scan_map
 from pytensor.scan import scan
@@ -17,16 +18,20 @@
 from pymc_experimental.distributions import DiscreteMarkovChain
 
 
-class MarginalRV(SymbolicRandomVariable):
+class MarginalRV(OpFromGraph, MeasurableOp):
     """Base class for Marginalized RVs"""
 
+    def __init__(self, *args, dims_connections: tuple[tuple[int | None]], **kwargs) -> None:
+        self.dims_connections = dims_connections
+        super().__init__(*args, **kwargs)
 
-class FiniteDiscreteMarginalRV(MarginalRV):
-    """Base class for Finite Discrete Marginalized RVs"""
 
+class MarginalFiniteDiscreteRV(MarginalRV):
+    """Base class for Marginalized Finite Discrete RVs"""
 
-class DiscreteMarginalMarkovChainRV(MarginalRV):
-    """Base class for Discrete Marginal Markov Chain RVs"""
+
+class MarginalDiscreteMarkovChainRV(MarginalRV):
+    """Base class for Marginalized Discrete Markov Chain RVs"""
 
 
 def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
@@ -48,24 +53,69 @@ def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
     raise NotImplementedError(f"Cannot compute domain for op {op}")
 
 
-def _add_reduce_batch_dependent_logps(
-    marginalized_type: TensorType, dependent_logps: Sequence[TensorVariable]
-):
-    """Add the logps of dependent RVs while reducing extra batch dims relative to `marginalized_type`."""
+def _reduce_batch_dependent_logps(
+    marginalized_op: MarginalRV,
+    marginalized_logp: TensorVariable,
+    dependent_ops: Sequence[Op],
+    dependent_logps: Sequence[TensorVariable],
+) -> TensorVariable:
+    """Combine the logps of dependent RVs with the marginalized logp.
+
+     This requires reducing extra batch dims and transposing when they are not aligned.
+
+        idx = pm.Bernoulli(idx, shape=(3, 2))  # 0, 1
+        pm.Normal("dep1", mu=idx.T[..., None] * 2, shape=(3, 2, 5))
+        pm.Normal("dep2", mu=idx * 2, shape=(7, 2, 3))
+
+        marginalize(idx)
+        dims_connections = [(1, 0, None), (None, 0, 1)]
+     """
+
+    dims_connections = marginalized_op.dims_connections
+
+    reduced_logps = [marginalized_logp]
+    for dependent_op, dependent_logp, dims_connection in zip(dependent_ops, dependent_logps, dims_connections):
+        if dependent_logp.type.ndim > 0:
+            # Find which support axis implied by the MarginalRV need to be reduced
+            # Some may have already been reduced by the logp expression of the dependent RV, for non-univariate RVs
+            if isinstance(dependent_op, MarginalRV):
+                dep_dims_connection = dependent_op.dims_connections[0]
+                dep_supp_axes = {-i for i, dim in enumerate(reversed(dep_dims_connection), start=1) if dim == ()}
+            else:
+                # For vanilla RVs, the support axes are the last ndim_supp
+                dep_supp_axes = set(range(-dependent_op.ndim_supp, 0))
+
+            # Dependent RV support axes are already collapsed in the logp, so we ignore them
+            supp_axes = [
+                -i
+                for i, dim in enumerate(reversed(dims_connection), start=1)
+                if (dim == () and -i not in dep_supp_axes)
+            ]
+
+            dependent_logp = dependent_logp.sum(supp_axes)
+            assert dependent_logp.type.ndim == marginalized_logp.type.ndim
 
-    mbcast = marginalized_type.broadcastable
-    reduced_logps = []
-    for dependent_logp in dependent_logps:
-        dbcast = dependent_logp.type.broadcastable
-        dim_diff = len(dbcast) - len(mbcast)
-        mbcast_aligned = mbcast + (True,) * dim_diff
-        vbcast_axis = [i for i, (m, v) in enumerate(zip(mbcast_aligned, dbcast)) if m and not v]
-        reduced_logps.append(dependent_logp.sum(vbcast_axis))
-    return pt.add(*reduced_logps)
+            # Finally, we need to align the dependent logp batch dimensions with the marginalized logp
+            dims_alignment = [dim[0] for dim in dims_connection if dim != ()]
+            dependent_logp = dependent_logp.transpose(*dims_alignment)
 
+        reduced_logps.append(dependent_logp)
 
-@_logprob.register(FiniteDiscreteMarginalRV)
-def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
+    reduced_logp = pt.add(*reduced_logps)
+
+    if reduced_logp.type.ndim > 0:
+        # Transpose reduced logp into the direction of the first dependent RV
+        first_dims_alignment = [dim[0] for dim in dims_connections[0] if dim != ()]
+        reduced_logp = reduced_logp.transpose(*first_dims_alignment)
+
+    return reduced_logp
+
+
+dummy_zero = pt.constant(0, name="dummy_zero")
+
+
+@_logprob.register(MarginalFiniteDiscreteRV)
+def finite_discrete_marginal_rv_logp(op: MarginalFiniteDiscreteRV, values, *inputs, **kwargs):
     # Clone the inner RV graph of the Marginalized RV
     marginalized_rvs_node = op.make_node(*inputs)
     marginalized_rv, *inner_rvs = clone_replace(
@@ -81,8 +131,11 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
 
     # Reduce logp dimensions corresponding to broadcasted variables
     marginalized_logp = logps_dict.pop(marginalized_vv)
-    joint_logp = marginalized_logp + _add_reduce_batch_dependent_logps(
-        marginalized_rv.type, logps_dict.values()
+    joint_logp = _reduce_batch_dependent_logps(
+        marginalized_op=op,
+        marginalized_logp=marginalized_logp,
+        dependent_ops=[inner_rv.owner.op for inner_rv in inner_rvs],
+        dependent_logps=[logps_dict[value] for value in values],
     )
 
     # Compute the joint_logp for all possible n values of the marginalized RV. We assume
@@ -119,21 +172,20 @@ def logp_fn(marginalized_rv_const, *non_sequences):
             mode=Mode().including("local_remove_check_parameter"),
         )
 
-    joint_logps = pt.logsumexp(joint_logps, axis=0)
+    joint_logp = pt.logsumexp(joint_logps, axis=0)
 
     # We have to add dummy logps for the remaining value variables, otherwise PyMC will raise
-    return joint_logps, *(pt.constant(0),) * (len(values) - 1)
+    return joint_logp, *((dummy_zero,) * (len(values) - 1))
 
 
-@_logprob.register(DiscreteMarginalMarkovChainRV)
+@_logprob.register(MarginalDiscreteMarkovChainRV)
 def marginal_hmm_logp(op, values, *inputs, **kwargs):
     marginalized_rvs_node = op.make_node(*inputs)
-    inner_rvs = clone_replace(
+    chain_rv, *dependent_rvs = clone_replace(
         op.inner_outputs,
         replace={u: v for u, v in zip(op.inner_inputs, marginalized_rvs_node.inputs)},
     )
 
-    chain_rv, *dependent_rvs = inner_rvs
     P, n_steps_, init_dist_, rng = chain_rv.owner.inputs
     domain = pt.arange(P.shape[-1], dtype="int32")
 
@@ -149,9 +201,11 @@ def marginal_hmm_logp(op, values, *inputs, **kwargs):
 
     # Reduce and add the batch dims beyond the chain dimension
     reduced_logp_emissions = _add_reduce_batch_dependent_logps(
+        init_logp,
         chain_rv.type, logp_emissions_dict.values()
     )
 
+
     # Add a batch dimension for the domain of the chain
     chain_shape = constant_fold(tuple(chain_rv.shape))
     batch_chain_value = pt.moveaxis(pt.full((*chain_shape, domain.size), domain), -1, 0)
@@ -188,7 +242,9 @@ def step_alpha(logp_emission, log_alpha, log_P):
     # Final logp is just the sum of the last scan state
     joint_logp = pt.logsumexp(log_alpha_seq[-1], axis=0)
 
+    # TODO: Transpose into shape of first emission
+
     # If there are multiple emission streams, we have to add dummy logps for the remaining value variables. The first
-    # return is the joint probability of everything together, but PyMC still expects one logp for each one.
-    dummy_logps = (pt.constant(0),) * (len(values) - 1)
+    # return is the joint probability of everything together, but PyMC still expects one logp for each emission stream.
+    dummy_logps = (dummy_zero) * (len(values) - 1)
     return joint_logp, *dummy_logps

pymc_experimental/model/marginal/graph_analysis.py

@@ -15,6 +15,8 @@
 from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor, get_idx_list
 from pytensor.tensor.type_other import NoneTypeT
 
+from pymc_experimental.model.marginal.distributions import MarginalRV
+
 
 def static_shape_ancestors(vars):
     """Identify ancestors Shape Ops of static shapes (therefore constant in a valid graph)."""
@@ -101,6 +103,12 @@ def _broadcast_dims(
     output_dims = tuple(
         tuple(sorted(set(chain.from_iterable(inputs_dim)))) for inputs_dim in zip(*inputs_dims)
     )
+    if any(len(output_dim) > 1 for output_dim in output_dims):
+        raise ValueError("Different known dimensions mixed via broadcasting")
+
+    if len(set(output_dim[0] for output_dim in output_dims if output_dim != ())) < len([output_dim for output_dim in output_dims if output_dim != ()]):
+        raise ValueError("Same dimension used in different axis after broadcasting")
+
     return output_dims
 
 
@@ -111,6 +119,9 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
     for node in io_toposort(input_vars, output_vars):
         inputs_dims = [var_dims.get(inp, ()) for inp in node.inputs]
 
+        # f(marginalized_rv, *other_junk) -> dep_rv1, dep_rv2
+        # g(marginalized_rv.ravel()[i], *other_junk) -> dep_rv1.ravel()[?], derp_rv2.ravel()[?]
+
         if not any(inputs_dims):
             # None of the inputs are related to the batch_axes of the marginalized_rv
             continue
@@ -122,10 +133,23 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
             )
             var_dims[node.outputs[0]] = output_dims
 
+        elif isinstance(node.op, MarginalRV):
+
+            inner_var_dims = {
+                inner_inp: input_dims
+                for inner_inp, input_dims in zip(node.op.inner_inputs, inputs_dims)
+            }
+            inner_var_dims = _subgraph_dim_connection(
+                inner_var_dims, node.op.inner_inputs, node.op.inner_outputs
+            )
+            for out, inner_out in zip(node.outputs, node.op.inner_outputs):
+                # FIXME: If the known output_dim belongs to the supp_axis of the inner MarginalizedRV, this should raise
+                # Add test in test_graph_analysis
+                if inner_out in inner_var_dims:
+                    var_dims[out] = inner_var_dims[inner_out]
+
         elif (
-            isinstance(node.op, CustomSymbolicDistRV)
-            or isinstance(node.op, SymbolicRandomVariable)
-            and node.op.extended_signature is None
+            isinstance(node.op, SymbolicRandomVariable) and node.op.extended_signature is None
         ):
             # SymbolicRandomVariables without signature are a wild-card, so we need to introspect the inner graph.
             # MarginalRVs are such a case!
@@ -274,6 +298,13 @@ def _subgraph_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR
 
             var_dims[node.outputs[0]] = output_dims
 
+        # categorical(p=dimshuffle(matrix))
+        # (0, 1) -> (1, 0) -> (1,)
+        # ((0,), (1,)) -> (0, 1)
+        # (a,b),(c,d)->(a,b,c,d)
+        # (),()->() -> (a, b),(None, None),(a * None,b * None)
+        # (a, b),(c, d),(a * c, b * d)
+
         else:
             raise NotImplementedError(f"Marginalization through operation {node} not supported.")
 

pymc_experimental/model/marginal/marginal_model.py

@@ -24,9 +24,9 @@
 
 from pymc_experimental.distributions import DiscreteMarkovChain
 from pymc_experimental.model.marginal.distributions import (
-    DiscreteMarginalMarkovChainRV,
-    FiniteDiscreteMarginalRV,
-    _add_reduce_batch_dependent_logps,
+    MarginalDiscreteMarkovChainRV,
+    MarginalFiniteDiscreteRV,
+    _reduce_batch_dependent_logps,
     get_domain_of_finite_discrete_rv,
 )
 from pymc_experimental.model.marginal.graph_analysis import (
@@ -431,7 +431,7 @@ def transform_input(inputs):
 
             # Handle batch dims for marginalized value and its dependent RVs
             marginalized_logp, *dependent_logps = joint_logps
-            joint_logp = marginalized_logp + _add_reduce_batch_dependent_logps(
+            joint_logp = marginalized_logp + _reduce_batch_dependent_logps(
                 marginalized_rv.type, dependent_logps
             )
 
@@ -556,12 +556,7 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
         if rv is not rv_to_marginalize
     ]
 
-    dependent_rvs_ndim_supp = {dependent_rv.owner.op.ndim_supp for dependent_rv in dependent_rvs}
-    if len(dependent_rvs_ndim_supp) > 1:
-        raise NotImplementedError("All dependent RVs must have the same support dimensionality")
-
-    [dependent_rv_ndim_supp] = dependent_rvs_ndim_supp
-
+    # Todo: back to broadcastable
     if rv_to_marginalize.type.ndim > 0:
         # If the marginalized RV has multiple dimensions, check that graph between
         # marginalized RV and dependent RVs does not mix information from batch dimensions
@@ -573,33 +568,12 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
         except ValueError as e:
             # For the perspective of the user this is a NotImplementedError
             raise NotImplementedError(
-                "The graph between the marginalized and dependent RVs cannot be marginalized"
+                "The graph between the marginalized and dependent RVs cannot be marginalized."
             ) from e
 
-        if any(
-            len(dim) > 1
-            for rv_dim_connections in dependent_rvs_dim_connections
-            for dim in rv_dim_connections
-        ):
-            raise NotImplementedError("Multiple dimensions are mixed")
-
-        # We further check that batch dimensions of the marginalized RVs are aligned with those of the dependent RV
-        marginal_ndim = rv_to_marginalize.type.ndim
-        marginal_batch_dims = tuple((i,) for i in range(marginal_ndim))
-        for dependent_rv, dependent_rv_batch_dims in zip(
-            dependent_rvs, dependent_rvs_dim_connections
-        ):
-            extra_batch_ndim = dependent_rv.type.ndim + dependent_rv_ndim_supp - marginal_ndim
-            valid_dependent_batch_dims = marginal_batch_dims + (((),) * extra_batch_ndim)
-            if dependent_rv_batch_dims != valid_dependent_batch_dims:
-                raise NotImplementedError(
-                    f"Link between dimensions of marginalized and dependent RVs not supported: {dependent_rv_batch_dims} != {valid_dependent_batch_dims}"
-                )
+    else:
+        dependent_rvs_dim_connections = tuple(((),) * dependent_rv.type.ndim for dependent_rv in dependent_rvs)
 
-    ndim_supp = max(
-        (dependent_rv.type.ndim + dependent_rv_ndim_supp - rv_to_marginalize.type.ndim)
-        for dependent_rv in dependent_rvs
-    )
 
     input_rvs = list(set((*marginalized_rv_input_rvs, *other_direct_rv_ancestors)))
     output_rvs = [rv_to_marginalize, *dependent_rvs]
@@ -608,14 +582,14 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     inputs = input_rvs + collect_shared_vars(output_rvs, blockers=input_rvs)
 
     if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
-        marginalize_constructor = DiscreteMarginalMarkovChainRV
+        marginalize_constructor = MarginalDiscreteMarkovChainRV
     else:
-        marginalize_constructor = FiniteDiscreteMarginalRV
+        marginalize_constructor = MarginalFiniteDiscreteRV
 
     marginalization_op = marginalize_constructor(
         inputs=inputs,
         outputs=output_rvs,  # TODO: Add RNG updates to outputs
-        ndim_supp=ndim_supp,
+        dims_connections=dependent_rvs_dim_connections,
     )
     new_output_rvs = marginalization_op(*inputs)
     fgraph.replace_all(tuple(zip(output_rvs, new_output_rvs)))

tests/model/marginal/test_distributions.py

@@ -8,7 +8,7 @@
 from pymc_experimental import MarginalModel
 from pymc_experimental.distributions import DiscreteMarkovChain
 
-from pymc_experimental.model.marginal.distributions import FiniteDiscreteMarginalRV
+from pymc_experimental.model.marginal.distributions import MarginalFiniteDiscreteRV
 
 
 def test_marginalized_bernoulli_logp():
@@ -17,13 +17,10 @@ def test_marginalized_bernoulli_logp():
 
     idx = pm.Bernoulli.dist(0.7, name="idx")
     y = pm.Normal.dist(mu=mu[idx], sigma=1.0, name="y")
-    marginal_rv_node = FiniteDiscreteMarginalRV(
+    marginal_rv_node = MarginalFiniteDiscreteRV(
         [mu],
         [idx, y],
-        ndim_supp=0,
-        n_updates=0,
-        # Ignore the fact we didn't specify shared RNG input/outputs for idx,y
-        strict=False,
+        dims_connections=(((),),),
     )(mu)[0].owner
 
     y_vv = y.clone()