WIP Automatic marginalization finite discrete variables

Author: ricardoV94

notebooks/marginalized_changepoint_model.ipynb

@@ -0,0 +1,361 @@
+from typing import Dict, Sequence, Tuple, Union
+
+import aesara.tensor as at
+import numpy as np
+from aeppl import factorized_joint_logprob
+from aeppl.abstract import _get_measurable_outputs
+from aeppl.logprob import _logprob
+from aesara import clone_replace
+from aesara.compile import SharedVariable
+from aesara.compile.builders import OpFromGraph
+from aesara.graph import Constant, FunctionGraph, ancestors
+from aesara.tensor import TensorVariable
+from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.random.var import (
+    RandomGeneratorSharedVariable,
+    RandomStateSharedVariable,
+)
+from pymc import SymbolicRandomVariable
+from pymc.aesaraf import constant_fold, inputvars
+from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
+from pymc.distributions.distribution import _moment, moment
+from pymc.model import Model
+
+
+class MarginalModel(Model):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.parent is not None:
+            raise NotImplementedError("MarginalModel cannot be used inside another Model")
+        else:
+            self.marginalized_rvs_to_dependent_rvs = {}
+
+    def logp(self, vars=None, **kwargs):
+        if not kwargs.get("sum", True):
+            # Check if dependent RVs were requested
+            if vars is not None and not isinstance(vars, Sequence):
+                vars = (vars,)
+            if vars is None or (
+                {v for vs in self.marginalized_rvs_to_dependent_rvs.values() for v in vs}
+                & {self.values_to_rvs.get(var, var) for var in vars}
+            ):
+                raise ValueError(
+                    "Cannot request elemwise logp (sum=False) for variables that depend on a marginalized RV"
+                )
+        return super().logp(vars, **kwargs)
+
+    def point_logps(self, *args, **kwargs):
+        # TODO: Fix this
+        return {}
+
+    def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorVariable]]):
+        # TODO: this does not need to be a property of a Model
+        if not isinstance(rvs_to_marginalize, Sequence):
+            rvs_to_marginalize = (rvs_to_marginalize,)
+
+        supported_dists = (Bernoulli, Categorical, DiscreteUniform)
+        for rv_to_marginalize in rvs_to_marginalize:
+            if rv_to_marginalize not in self.free_RVs:
+                raise ValueError(
+                    f"Marginalized RV {rv_to_marginalize} is not a free RV in the model"
+                )
+            if not isinstance(rv_to_marginalize.owner.op, supported_dists):
+                raise NotImplementedError(
+                    f"RV with distribution {rv_to_marginalize.owner.op} cannot be marginalized. "
+                    f"Supported distribution include {supported_dists}"
+                )
+
+        if self.deterministics:
+            # TODO: This should be fine if deterministics do not depend on marginalized RVs
+            raise NotImplementedError("Models with deterministics cannot be marginalized")
+
+        if self.potentials:
+            raise NotImplementedError("Models with potentials cannot be marginalized")
+
+        # Replaced with subgraph that need to be marginalized for each RV
+        fg = FunctionGraph(outputs=self.basic_RVs, clone=False)
+        toposort = fg.toposort()
+        replacements = {}
+        new_marginalized_rv = None
+        new_dependent_rvs = []
+        for rv_to_marginalize in sorted(
+            rvs_to_marginalize, key=lambda rv: toposort.index(rv.owner)
+        ):
+            old_rvs, new_rvs = _replace_finite_discrete_marginal_subgraph(
+                fg, rv_to_marginalize, self.rvs_to_values
+            )
+            # Update old mappings
+            for old_rv, new_rv in zip(old_rvs, new_rvs):
+                replacements[old_rv] = new_rv
+
+                value = self.rvs_to_values.pop(old_rv)
+                self.named_vars.pop(old_rv.name)
+                new_rv.name = old_rv.name
+
+                if old_rv is rv_to_marginalize:
+                    self.free_RVs.remove(old_rv)
+                    self.values_to_rvs.pop(value)
+                    self.rvs_to_transforms.pop(old_rv)
+                    self.rvs_to_total_sizes.pop(old_rv)
+                    new_marginalized_rv = new_rv
+                    continue
+
+                new_dependent_rvs.append(new_rv)
+                if old_rv in self.free_RVs:
+                    index = self.free_RVs.index(old_rv)
+                    self.free_RVs.pop(index)
+                    self.free_RVs.insert(index, new_rv)
+                    self._initial_values[new_rv] = self._initial_values.pop(old_rv)
+                else:
+                    index = self.observed_RVs.index(old_rv)
+                    self.observed_RVs.pop(index)
+                    self.observed_RVs.insert(index, new_rv)
+                self.rvs_to_values[new_rv] = value
+                self.named_vars[new_rv.name] = new_rv
+                self.values_to_rvs[value] = new_rv
+                self.rvs_to_transforms[new_rv] = self.rvs_to_transforms.pop(old_rv)
+                # TODO: Automatic imputation RV does not seem to have total_size mapping
+                self.rvs_to_total_sizes[new_rv] = self.rvs_to_total_sizes.pop(old_rv, None)
+
+        self.marginalized_rvs_to_dependent_rvs[new_marginalized_rv] = new_dependent_rvs
+        return replacements
+
+
+def _find_dependent_rvs(dependable_rv, all_rvs):
+    # Find rvs than depend on dependable
+    dependent_rvs = []
+    for rv in all_rvs:
+        if rv is dependable_rv:
+            continue
+        blockers = [other_rv for other_rv in all_rvs if other_rv is not rv]
+        if dependable_rv in ancestors([rv], blockers=blockers):
+            dependent_rvs.append(rv)
+    return dependent_rvs
+
+
+def _find_input_rvs(output_rvs, all_rvs):
+    blockers = [other_rv for other_rv in all_rvs if other_rv not in 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)))
+    ]
+
+
+def _is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs):
+    # TODO: No need to consider apply nodes outside the subgraph...
+    fg = FunctionGraph(outputs=output_rvs, clone=False)
+
+    non_elemwise_blockers = [
+        o for node in fg.apply_nodes if not isinstance(node.op, Elemwise) for o in node.outputs
+    ]
+    blocker_candidates = [rv_to_marginalize] + other_input_rvs + non_elemwise_blockers
+    blockers = [var for var in blocker_candidates if var not in output_rvs]
+
+    # TODO: We could actually use these truncated inputs to
+    # generate a smaller Marginalized graph...
+    truncated_inputs = [
+        var
+        for var in ancestors(output_rvs, blockers=blockers)
+        if (
+            var in blockers
+            or (var.owner is None and not isinstance(var, (Constant, SharedVariable)))
+        )
+    ]
+
+    # Check that we reach the marginalized rv following a pure elemwise graph
+    if rv_to_marginalize not in truncated_inputs:
+        return False
+
+    # Check that none of the truncated inputs depends on the marginalized_rv
+    other_truncated_inputs = [inp for inp in truncated_inputs if inp is not rv_to_marginalize]
+    # TODO: We don't need to go all the way to the root variables
+    if rv_to_marginalize in ancestors(
+        other_truncated_inputs, blockers=[rv_to_marginalize, *other_input_rvs]
+    ):
+        return False
+    return True
+
+
+SUPPORTED_RNG_TYPES = (RandomStateSharedVariable, RandomGeneratorSharedVariable)
+
+
+class FiniteDiscreteMarginalRV(SymbolicRandomVariable):
+    def __init__(self, *args, n_updates: int, **kwargs):
+        self.n_updates = n_updates
+        super().__init__(*args, **kwargs)
+
+    def update(self, node):
+        n_updates = node.op.n_updates
+        shared_rng_inputs = node.inputs[:n_updates]
+        update_outputs = node.outputs[:n_updates]
+        assert len(update_outputs) == len(shared_rng_inputs)
+        # We made sure to pass RNG inputs and output updates in the same order
+        return {inp: out for inp, out in zip(shared_rng_inputs, update_outputs)}
+
+
+def _collect_updates(rvs: Sequence[TensorVariable]) -> Dict[TensorVariable, TensorVariable]:
+    rng_updates = {}
+    for rv in rvs:
+        if isinstance(rv.owner.op, RandomVariable):
+            rng = rv.owner.inputs[0]
+            assert not hasattr(rng, "default_update")
+            rng_updates[rng] = rv.owner.outputs[0]
+        elif isinstance(rv.owner.op, SymbolicRandomVariable):
+            rng_updates.update(rv.owner.op.udpate(rv.owner))
+        else:
+            raise TypeError(f"Unknown RV type: {rv.owner.op}")
+    assert all(isinstance(rng, SUPPORTED_RNG_TYPES) for rng in rng_updates.keys())
+    return rng_updates
+
+
+def _replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, rvs_to_values):
+    # TODO: This should eventually be integrated in a more general routine that can
+    #  identify other types of supported marginalization, of which finite discrete
+    #  RVs is just one
+
+    dependent_rvs = _find_dependent_rvs(rv_to_marginalize, rvs_to_values)
+    if not dependent_rvs:
+        raise ValueError(f"No RVs depend on marginalized RV {rv_to_marginalize}")
+
+    marginalized_rv_input_rvs = _find_input_rvs([rv_to_marginalize], rvs_to_values)
+    dependent_rvs_input_rvs = [
+        rv for rv in _find_input_rvs(dependent_rvs, rvs_to_values) if rv is not rv_to_marginalize
+    ]
+
+    # If the marginalized RV has batched dimensions, check that graph between
+    # marginalized RV and dependent RVs is composed strictly of Elemwise Operations.
+    # This implies (?) that the dimensions are completely independent and a logp graph
+    # can ultimately be generated that is proportional to the support domain and not
+    # We don't need to worry about batched graphs if the  RV is scalar.
+    # TODO: This eval is a bit hackish
+    if np.prod(rv_to_marginalize.shape.eval()) > 1:
+        if not _is_elemwise_subgraph(rv_to_marginalize, dependent_rvs_input_rvs, dependent_rvs):
+            raise NotImplementedError(
+                "The subgraph between a marginalized RV and its dependents includes non Elemwise operations. "
+                "This is currently not supported",
+            )
+
+    input_rvs = [*marginalized_rv_input_rvs, *dependent_rvs_input_rvs]
+    rvs_to_marginalize = [rv_to_marginalize, *dependent_rvs]
+
+    # Collect update expressions of the inner RVs.
+    # Note: This could be avoided if we inlined the MarginalOp Graph before collecting
+    # the updates in `pymc.aesaraf.compile_pymc`
+    updates_rvs_to_marginalize = _collect_updates(rvs_to_marginalize)
+    n_updates = len(updates_rvs_to_marginalize)
+    assert n_updates
+
+    outputs = list(updates_rvs_to_marginalize.values()) + rvs_to_marginalize
+    # Clone replace inner RV rng inputs so that we can be sure of the update order
+    replace_inputs = {rng: rng.type() for rng in updates_rvs_to_marginalize.keys()}
+    # Clone replace outter RV inputs, so that their shared RNGs don't make it into
+    # the inner graph of the marginalized RVs
+    replace_inputs.update({input_rv: input_rv.type() for input_rv in input_rvs})
+    cloned_outputs = clone_replace(outputs, replace=replace_inputs)
+
+    marginalization_op = FiniteDiscreteMarginalRV(
+        inputs=list(replace_inputs.values()),
+        outputs=cloned_outputs,
+        ndim_supp=-1,  # This will certainly break stuff :D
+        n_updates=n_updates,
+    )
+    marginalized_rvs = marginalization_op(*replace_inputs.keys())[n_updates:]
+    fgraph.replace_all(tuple(zip(rvs_to_marginalize, marginalized_rvs)))
+    return rvs_to_marginalize, marginalized_rvs
+
+
+@_get_measurable_outputs.register(FiniteDiscreteMarginalRV)
+def _get_measurable_outputs_finite_discrete_marginal_rv(op, node):
+    # The Marginalized RV (first non-update output) is not measurable, nor are updates
+    return node.outputs[op.n_updates + 1 :]
+
+
+@_moment.register(FiniteDiscreteMarginalRV)
+def moment_finite_discrete_marginal_rv(op, rv, *rv_inputs):
+    # Recreate inner RV and retrieve its moment
+    node = rv.owner
+    marginalized_rv, *dependent_rvs = clone_replace(
+        op.inner_outputs[op.n_updates :],
+        replace={u: v for u, v in zip(op.inner_inputs, rv_inputs)},
+    )
+    rv_idx = node.outputs[op.n_updates + 1 :].index(rv)
+    rv = dependent_rvs[rv_idx]
+
+    moment_marginalized_rv = moment(marginalized_rv)
+    (rv,) = clone_replace([rv], replace={marginalized_rv: moment_marginalized_rv})
+    return moment(rv)
+
+
+def _get_domain_of_finite_discrete_rv(rv: TensorVariable) -> Tuple[int, ...]:
+    op = rv.owner.op
+    if isinstance(op, Bernoulli):
+        return (0, 1)
+    elif isinstance(op, Categorical):
+        p_param = rv.owner.inputs[3]
+        return tuple(range(at.get_vector_length(p_param)))
+    elif isinstance(op, DiscreteUniform):
+        lower, upper = constant_fold(rv.owner.inputs[3:])
+        return tuple(range(lower, upper + 1))
+
+    raise NotImplementedError(f"Cannot compute domain for op {op}")
+
+
+@_logprob.register(FiniteDiscreteMarginalRV)
+def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
+
+    marginalized_rvs_node = op.make_node(*inputs)
+    marginalized_rv, *dependent_rvs = clone_replace(
+        op.inner_outputs[op.n_updates :],
+        replace={u: v for u, v in zip(op.inner_inputs, marginalized_rvs_node.inputs)},
+    )
+
+    # Some inputs are not root inputs (such as transformed projections of value variables)
+    # Or cannot be used as inputs to an OpFromGraph (shared variables and constants)
+    inputs = list(inputvars(inputs))
+
+    rvs_to_values = {}
+    dummy_marginalized_value = marginalized_rv.clone()
+    rvs_to_values[marginalized_rv] = dummy_marginalized_value
+    rvs_to_values.update(zip(dependent_rvs, values))
+    _logp = at.sum(
+        [
+            at.sum(factor)
+            for factor in factorized_joint_logprob(rv_values=rvs_to_values, **kwargs).values()
+        ]
+    )
+    # OpFromGraph does not accept constant inputs...
+    _values = [
+        value
+        for value in rvs_to_values.values()
+        if not isinstance(value, (Constant, SharedVariable))
+    ]
+    # TODO: If we inline the logp graph, optimization becomes incredibly painful for
+    #  large domains... Would be great to find a way to vectorize the graph across
+    #  the domain values (when possible)
+    logp_op = OpFromGraph([*_values, *inputs], [_logp], inline=False)
+
+    # PyMC does not allow RVs in the logp graph... Even if we are just using the shape
+    # TODO: Get better work-around that .eval(). It probably makes sense to do a constant
+    #  fold pass in the final logp graph, so that individual logp functions don't have
+    #  to worry about it
+    marginalized_rv_shape = marginalized_rv.shape.eval()
+    non_const_values = [
+        value for value in values if not isinstance(value, (Constant, SharedVariable))
+    ]
+    logp = at.logsumexp(
+        [
+            logp_op(
+                np.full(marginalized_rv_shape, marginalized_rv_const), *non_const_values, *inputs
+            )
+            for marginalized_rv_const in _get_domain_of_finite_discrete_rv(marginalized_rv)
+        ]
+    )
+    # In the case of multiple dependent values, the whole logp is assigned just to the
+    # first value. This is a quite hackish, but Aeppl errors out if some value variable
+    # is not assigned a specific logp term, and it also does not make sense to separate
+    # them internally.
+    dummy_logps = (at.constant([], name="dummy_marginalized_logp"),) * (len(values) - 1)
+    return logp, *dummy_logps