pymc-extras 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

pymc_extras/model/marginal/marginal_model.py

@@ -1,7 +1,6 @@
 import warnings
 
 from collections.abc import Sequence
-from typing import Union
 
 import numpy as np
 import pymc
@@ -13,21 +12,42 @@ from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.transforms import Chain
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
-from pymc.pytensorf import compile_pymc, constant_fold
-from pymc.util import RandomState, _get_seeds_per_chain, treedict
+from pymc.model.fgraph import (
+    ModelFreeRV,
+    ModelValuedVar,
+    fgraph_from_model,
+    model_free_rv,
+    model_from_fgraph,
+)
+from pymc.pytensorf import collect_default_updates, compile_pymc, constant_fold, toposort_replace
+from pymc.util import RandomState, _get_seeds_per_chain
+from pytensor import In, Out
 from pytensor.compile import SharedVariable
-from pytensor.graph import FunctionGraph, clone_replace, graph_inputs
-from pytensor.graph.replace import vectorize_graph
+from pytensor.graph import (
+    FunctionGraph,
+    Variable,
+    clone_replace,
+    graph_inputs,
+    graph_replace,
+    node_rewriter,
+    vectorize_graph,
+)
+from pytensor.graph.rewriting.basic import in2out
 from pytensor.tensor import TensorVariable
-from pytensor.tensor.special import log_softmax
 
 __all__ = ["MarginalModel", "marginalize"]
 
+from pytensor.tensor.random.type import RandomType
+from pytensor.tensor.special import log_softmax
+
 from pymc_extras.distributions import DiscreteMarkovChain
 from pymc_extras.model.marginal.distributions import (
     MarginalDiscreteMarkovChainRV,
     MarginalFiniteDiscreteRV,
+    MarginalRV,
+    NonSeparableLogpWarning,
     get_domain_of_finite_discrete_rv,
+    inline_ofg_outputs,
     reduce_batch_dependent_logps,
 )
 from pymc_extras.model.marginal.graph_analysis import (
@@ -87,479 +107,452 @@ class MarginalModel(Model):
     """
 
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.marginalized_rvs = []
-        self._marginalized_named_vars_to_dims = {}
+        raise TypeError(
+            "MarginalModel was deprecated in favor of `marginalize` which now returns a PyMC model"
+        )
 
-    def _delete_rv_mappings(self, rv: TensorVariable) -> None:
-        """Remove all model mappings referring to rv
 
-        This can be used to "delete" an RV from a model
-        """
-        assert rv in self.basic_RVs, "rv is not part of the Model"
+def _warn_interval_transform(rv_to_marginalize, replaced_vars: Sequence[ModelValuedVar]) -> None:
+    for replaced_var in replaced_vars:
+        if not isinstance(replaced_var.owner.op, ModelValuedVar):
+            raise TypeError(f"{replaced_var} is not a ModelValuedVar")
 
-        name = rv.name
-        self.named_vars.pop(name)
-        if name in self.named_vars_to_dims:
-            self.named_vars_to_dims.pop(name)
+        if not isinstance(replaced_var.owner.op, ModelFreeRV):
+            continue
 
-        value = self.rvs_to_values.pop(rv)
-        self.values_to_rvs.pop(value)
+        if replaced_var is rv_to_marginalize:
+            continue
 
-        self.rvs_to_transforms.pop(rv)
-        if rv in self.free_RVs:
-            self.free_RVs.remove(rv)
-            self.rvs_to_initial_values.pop(rv)
-        else:
-            self.observed_RVs.remove(rv)
-
-    def _transfer_rv_mappings(self, old_rv: TensorVariable, new_rv: TensorVariable) -> None:
-        """Transfer model mappings from old_rv to new_rv"""
-
-        assert old_rv in self.basic_RVs, "old_rv is not part of the Model"
-        assert new_rv not in self.basic_RVs, "new_rv is already part of the Model"
-
-        self.named_vars.pop(old_rv.name)
-        new_rv.name = old_rv.name
-        self.named_vars[new_rv.name] = new_rv
-        if old_rv in self.named_vars_to_dims:
-            self._RV_dims[new_rv] = self._RV_dims.pop(old_rv)
-
-        value = self.rvs_to_values.pop(old_rv)
-        self.rvs_to_values[new_rv] = value
-        self.values_to_rvs[value] = new_rv
-
-        self.rvs_to_transforms[new_rv] = self.rvs_to_transforms.pop(old_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.rvs_to_initial_values[new_rv] = self.rvs_to_initial_values.pop(old_rv)
-        elif old_rv in self.observed_RVs:
-            index = self.observed_RVs.index(old_rv)
-            self.observed_RVs.pop(index)
-            self.observed_RVs.insert(index, new_rv)
-
-    def _marginalize(self, user_warnings=False):
-        fg = FunctionGraph(outputs=self.basic_RVs + self.marginalized_rvs, clone=False)
-
-        toposort = fg.toposort()
-        rvs_left_to_marginalize = self.marginalized_rvs
-        for rv_to_marginalize in sorted(
-            self.marginalized_rvs,
-            key=lambda rv: toposort.index(rv.owner),
-            reverse=True,
+        transform = replaced_var.owner.op.transform
+
+        if isinstance(transform, IntervalTransform) or (
+            isinstance(transform, Chain)
+            and any(isinstance(tr, IntervalTransform) for tr in transform.transform_list)
         ):
-            # Check that no deterministics or potentials dependend on the rv to marginalize
-            for det in self.deterministics:
-                if is_conditional_dependent(
-                    det, rv_to_marginalize, self.basic_RVs + rvs_left_to_marginalize
-                ):
-                    raise NotImplementedError(
-                        f"Cannot marginalize {rv_to_marginalize} due to dependent Deterministic {det}"
-                    )
-            for pot in self.potentials:
-                if is_conditional_dependent(
-                    pot, rv_to_marginalize, self.basic_RVs + rvs_left_to_marginalize
-                ):
-                    raise NotImplementedError(
-                        f"Cannot marginalize {rv_to_marginalize} due to dependent Potential {pot}"
-                    )
-
-            old_rvs, new_rvs = replace_finite_discrete_marginal_subgraph(
-                fg, rv_to_marginalize, self.basic_RVs + rvs_left_to_marginalize
+            warnings.warn(
+                f"The transform {transform} for the variable {replaced_var}, which depends on the "
+                f"marginalized {rv_to_marginalize} may no longer work if bounds depended on other variables.",
+                UserWarning,
             )
 
-            if user_warnings and len(new_rvs) > 2:
-                warnings.warn(
-                    "There are multiple dependent variables in a FiniteDiscreteMarginalRV. "
-                    f"Their joint logp terms will be assigned to the first RV: {old_rvs[1]}",
-                    UserWarning,
+
+def _unique(seq: Sequence) -> list:
+    """Copied from https://stackoverflow.com/a/480227"""
+    seen = set()
+    seen_add = seen.add
+    return [x for x in seq if not (x in seen or seen_add(x))]
+
+
+def marginalize(model: Model, rvs_to_marginalize: ModelRVs) -> MarginalModel:
+    """Marginalize a subset of variables in a PyMC model.
+
+    This creates a class of `MarginalModel` from an existing `Model`, with the specified
+    variables marginalized.
+
+    See documentation for `MarginalModel` for more information.
+
+    Parameters
+    ----------
+    model : Model
+        PyMC model to marginalize. Original variables well be cloned.
+    rvs_to_marginalize : Sequence[TensorVariable]
+        Variables to marginalize in the returned model.
+
+    Returns
+    -------
+    marginal_model: MarginalModel
+        Marginal model with the specified variables marginalized.
+    """
+    if isinstance(rvs_to_marginalize, str | Variable):
+        rvs_to_marginalize = (rvs_to_marginalize,)
+
+    rvs_to_marginalize = [model[rv] if isinstance(rv, str) else rv for rv in rvs_to_marginalize]
+
+    if not rvs_to_marginalize:
+        return model
+
+    for rv_to_marginalize in rvs_to_marginalize:
+        if rv_to_marginalize not in model.free_RVs:
+            raise ValueError(f"Marginalized RV {rv_to_marginalize} is not a free RV in the model")
+
+        rv_op = rv_to_marginalize.owner.op
+        if isinstance(rv_op, DiscreteMarkovChain):
+            if rv_op.n_lags > 1:
+                raise NotImplementedError(
+                    "Marginalization for DiscreteMarkovChain with n_lags > 1 is not supported"
+                )
+            if rv_to_marginalize.owner.inputs[0].type.ndim > 2:
+                raise NotImplementedError(
+                    "Marginalization for DiscreteMarkovChain with non-matrix transition probability is not supported"
                 )
+        elif not isinstance(rv_op, Bernoulli | Categorical | DiscreteUniform):
+            raise NotImplementedError(
+                f"Marginalization of RV with distribution {rv_to_marginalize.owner.op} is not supported"
+            )
 
-            rvs_left_to_marginalize.remove(rv_to_marginalize)
-
-            for old_rv, new_rv in zip(old_rvs, new_rvs):
-                new_rv.name = old_rv.name
-                if old_rv in self.marginalized_rvs:
-                    idx = self.marginalized_rvs.index(old_rv)
-                    self.marginalized_rvs.pop(idx)
-                    self.marginalized_rvs.insert(idx, new_rv)
-                if old_rv in self.basic_RVs:
-                    self._transfer_rv_mappings(old_rv, new_rv)
-                    if user_warnings:
-                        # Interval transforms for dependent variable won't work for non-constant bounds because
-                        # the RV inputs are now different and may depend on another RV that also depends on the
-                        # same marginalized RV
-                        transform = self.rvs_to_transforms[new_rv]
-                        if isinstance(transform, IntervalTransform) or (
-                            isinstance(transform, Chain)
-                            and any(
-                                isinstance(tr, IntervalTransform) for tr in transform.transform_list
-                            )
-                        ):
-                            warnings.warn(
-                                f"The transform {transform} for the variable {old_rv}, which depends on the "
-                                f"marginalized {rv_to_marginalize} may no longer work if bounds depended on other variables.",
-                                UserWarning,
-                            )
-        return self
-
-    def _logp(self, *args, **kwargs):
-        return super().logp(*args, **kwargs)
-
-    def logp(self, vars=None, **kwargs):
-        m = self.clone()._marginalize()
-        if vars is not None:
-            if not isinstance(vars, Sequence):
-                vars = (vars,)
-            vars = [m[var.name] for var in vars]
-        return m._logp(vars=vars, **kwargs)
-
-    @staticmethod
-    def from_model(model: Union[Model, "MarginalModel"]) -> "MarginalModel":
-        new_model = MarginalModel(coords=model.coords)
-        if isinstance(model, MarginalModel):
-            marginalized_rvs = model.marginalized_rvs
-            marginalized_named_vars_to_dims = model._marginalized_named_vars_to_dims
-        else:
-            marginalized_rvs = []
-            marginalized_named_vars_to_dims = {}
-
-        model_vars = model.basic_RVs + model.potentials + model.deterministics + marginalized_rvs
-        data_vars = [var for name, var in model.named_vars.items() if var not in model_vars]
-        vars = model_vars + data_vars
-        cloned_vars = clone_replace(vars)
-        vars_to_clone = {var: cloned_var for var, cloned_var in zip(vars, cloned_vars)}
-        new_model.vars_to_clone = vars_to_clone
-
-        new_model.named_vars = treedict(
-            {name: vars_to_clone[var] for name, var in model.named_vars.items()}
-        )
-        new_model.named_vars_to_dims = model.named_vars_to_dims
-        new_model.values_to_rvs = {vv: vars_to_clone[rv] for vv, rv in model.values_to_rvs.items()}
-        new_model.rvs_to_values = {vars_to_clone[rv]: vv for rv, vv in model.rvs_to_values.items()}
-        new_model.rvs_to_transforms = {
-            vars_to_clone[rv]: tr for rv, tr in model.rvs_to_transforms.items()
-        }
-        new_model.rvs_to_initial_values = {
-            vars_to_clone[rv]: iv for rv, iv in model.rvs_to_initial_values.items()
-        }
-        new_model.free_RVs = [vars_to_clone[rv] for rv in model.free_RVs]
-        new_model.observed_RVs = [vars_to_clone[rv] for rv in model.observed_RVs]
-        new_model.potentials = [vars_to_clone[pot] for pot in model.potentials]
-        new_model.deterministics = [vars_to_clone[det] for det in model.deterministics]
-
-        new_model.marginalized_rvs = [vars_to_clone[rv] for rv in marginalized_rvs]
-        new_model._marginalized_named_vars_to_dims = marginalized_named_vars_to_dims
-        return new_model
-
-    def clone(self):
-        return self.from_model(self)
-
-    def marginalize(
-        self,
-        rvs_to_marginalize: ModelRVs,
+    fg, memo = fgraph_from_model(model)
+    rvs_to_marginalize = [memo[rv] for rv in rvs_to_marginalize]
+    toposort = fg.toposort()
+
+    for rv_to_marginalize in sorted(
+        rvs_to_marginalize,
+        key=lambda rv: toposort.index(rv.owner),
+        reverse=True,
     ):
-        if not isinstance(rvs_to_marginalize, Sequence):
-            rvs_to_marginalize = (rvs_to_marginalize,)
-        rvs_to_marginalize = [
-            self[var] if isinstance(var, str) else var for var in rvs_to_marginalize
-        ]
+        all_rvs = [node.out for node in fg.toposort() if isinstance(node.op, ModelValuedVar)]
+
+        dependent_rvs = find_conditional_dependent_rvs(rv_to_marginalize, all_rvs)
+        if not dependent_rvs:
+            # TODO: This should at most be a warning, not an error
+            raise ValueError(f"No RVs depend on marginalized RV {rv_to_marginalize}")
+
+        # Issue warning for IntervalTransform on dependent RVs
+        for dependent_rv in dependent_rvs:
+            transform = dependent_rv.owner.op.transform
 
-        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 isinstance(transform, IntervalTransform) or (
+                isinstance(transform, Chain)
+                and any(isinstance(tr, IntervalTransform) for tr in transform.transform_list)
+            ):
+                warnings.warn(
+                    f"The transform {transform} for the variable {dependent_rv}, which depends on the "
+                    f"marginalized {rv_to_marginalize} may no longer work if bounds depended on other variables.",
+                    UserWarning,
                 )
 
-            rv_op = rv_to_marginalize.owner.op
-            if isinstance(rv_op, DiscreteMarkovChain):
-                if rv_op.n_lags > 1:
-                    raise NotImplementedError(
-                        "Marginalization for DiscreteMarkovChain with n_lags > 1 is not supported"
-                    )
-                if rv_to_marginalize.owner.inputs[0].type.ndim > 2:
-                    raise NotImplementedError(
-                        "Marginalization for DiscreteMarkovChain with non-matrix transition probability is not supported"
-                    )
-            elif not isinstance(rv_op, Bernoulli | Categorical | DiscreteUniform):
+        # Check that no deterministics or potentials depend on the rv to marginalize
+        for det in model.deterministics:
+            if is_conditional_dependent(memo[det], rv_to_marginalize, all_rvs):
                 raise NotImplementedError(
-                    f"Marginalization of RV with distribution {rv_to_marginalize.owner.op} is not supported"
+                    f"Cannot marginalize {rv_to_marginalize} due to dependent Deterministic {det}"
+                )
+        for pot in model.potentials:
+            if is_conditional_dependent(memo[pot], rv_to_marginalize, all_rvs):
+                raise NotImplementedError(
+                    f"Cannot marginalize {rv_to_marginalize} due to dependent Potential {pot}"
                 )
 
-            if rv_to_marginalize.name in self.named_vars_to_dims:
-                dims = self.named_vars_to_dims[rv_to_marginalize.name]
-                self._marginalized_named_vars_to_dims[rv_to_marginalize.name] = dims
-
-            self._delete_rv_mappings(rv_to_marginalize)
-            self.marginalized_rvs.append(rv_to_marginalize)
-
-        # Raise errors and warnings immediately
-        self.clone()._marginalize(user_warnings=True)
-
-    def _to_transformed(self):
-        """Create a function from the untransformed space to the transformed space"""
-        transformed_rvs = []
-        transformed_names = []
-
-        for rv in self.free_RVs:
-            transform = self.rvs_to_transforms.get(rv)
-            if transform is None:
-                transformed_rvs.append(rv)
-                transformed_names.append(rv.name)
-            else:
-                transformed_rv = transform.forward(rv, *rv.owner.inputs)
-                transformed_rvs.append(transformed_rv)
-                transformed_names.append(self.rvs_to_values[rv].name)
-
-        fn = self.compile_fn(inputs=self.free_RVs, outs=transformed_rvs)
-        return fn, transformed_names
-
-    def unmarginalize(self, rvs_to_unmarginalize: Sequence[TensorVariable | str]):
-        for rv in rvs_to_unmarginalize:
-            if isinstance(rv, str):
-                rv = self[rv]
-            self.marginalized_rvs.remove(rv)
-            if rv.name in self._marginalized_named_vars_to_dims:
-                dims = self._marginalized_named_vars_to_dims.pop(rv.name)
-            else:
-                dims = None
-            self.register_rv(rv, name=rv.name, dims=dims)
-
-    def recover_marginals(
-        self,
-        idata: InferenceData,
-        var_names: Sequence[str] | None = None,
-        return_samples: bool = True,
-        extend_inferencedata: bool = True,
-        random_seed: RandomState = None,
-    ):
-        """Computes posterior log-probabilities and samples of marginalized variables
-        conditioned on parameters of the model given InferenceData with posterior group
+        marginalized_rv_input_rvs = find_conditional_input_rvs([rv_to_marginalize], all_rvs)
+        other_direct_rv_ancestors = [
+            rv
+            for rv in find_conditional_input_rvs(dependent_rvs, all_rvs)
+            if rv is not rv_to_marginalize
+        ]
+        input_rvs = _unique((*marginalized_rv_input_rvs, *other_direct_rv_ancestors))
 
-        When there are multiple marginalized variables, each marginalized variable is
-        conditioned on both the parameters and the other variables still marginalized
+        replace_finite_discrete_marginal_subgraph(fg, rv_to_marginalize, dependent_rvs, input_rvs)
 
-        All log-probabilities are within the transformed space
+    return model_from_fgraph(fg, mutate_fgraph=True)
 
-        Parameters
-        ----------
-        idata : InferenceData
-            InferenceData with posterior group
-        var_names : sequence of str, optional
-            List of variable names for which to compute posterior log-probabilities and samples. Defaults to all marginalized variables
-        return_samples : bool, default True
-            If True, also return samples of the marginalized variables
-        extend_inferencedata : bool, default True
-            Whether to extend the original InferenceData or return a new one
-        random_seed: int, array-like of int or SeedSequence, optional
-            Seed used to generating samples
 
-        Returns
-        -------
-        idata : InferenceData
-            InferenceData with where a lp_{varname} and {varname} for each marginalized variable in var_names added to the posterior group
+@node_rewriter(tracks=[MarginalRV])
+def local_unmarginalize(fgraph, node):
+    unmarginalized_rv, *dependent_rvs_and_rngs = inline_ofg_outputs(node.op, node.inputs)
+    rngs = [rng for rng in dependent_rvs_and_rngs if isinstance(rng.type, RandomType)]
+    dependent_rvs = [rv for rv in dependent_rvs_and_rngs if rv not in rngs]
 
-        .. code-block:: python
+    # Wrap the marginalized RV in a FreeRV
+    # TODO: Preserve dims and transform in MarginalRV
+    value = unmarginalized_rv.clone()
+    fgraph.add_input(value)
+    transform = None
+    unmarginalized_free_rv = model_free_rv(unmarginalized_rv, value, transform, *node.op.dims)
 
-            import pymc as pm
-            from pymc_extras import MarginalModel
+    # Replace references to the marginalized RV with the FreeRV in the dependent RVs
+    dependent_rvs = graph_replace(dependent_rvs, {unmarginalized_rv: unmarginalized_free_rv})
 
-            with MarginalModel() as m:
-                p = pm.Beta("p", 1, 1)
-                x = pm.Bernoulli("x", p=p, shape=(3,))
-                y = pm.Normal("y", pm.math.switch(x, -10, 10), observed=[10, 10, -10])
+    return [unmarginalized_free_rv, *dependent_rvs, *rngs]
 
-                m.marginalize([x])
 
-                idata = pm.sample()
-                m.recover_marginals(idata, var_names=["x"])
+unmarginalize_rewrite = in2out(local_unmarginalize, ignore_newtrees=False)
 
 
-        """
-        if var_names is None:
-            var_names = [var.name for var in self.marginalized_rvs]
+def unmarginalize(model: Model, rvs_to_unmarginalize: str | Sequence[str] | None = None) -> Model:
+    """Unmarginalize a subset of variables in a PyMC model.
 
-        var_names = [var if isinstance(var, str) else var.name for var in var_names]
-        vars_to_recover = [v for v in self.marginalized_rvs if v.name in var_names]
-        missing_names = [v.name for v in vars_to_recover if v not in self.marginalized_rvs]
-        if missing_names:
-            raise ValueError(f"Unrecognized var_names: {missing_names}")
 
-        if return_samples and random_seed is not None:
-            seeds = _get_seeds_per_chain(random_seed, len(vars_to_recover))
+    Parameters
+    ----------
+    model : Model
+        PyMC model to unmarginalize. Original variables well be cloned.
+    rvs_to_unmarginalize : str or sequence of str, optional
+        Variables to unmarginalize in the returned model. If None, all variables are
+        unmarginalized.
+
+    Returns
+    -------
+    unmarginal_model: Model
+        Model with the specified variables unmarginalized.
+    """
+
+    # Unmarginalize all the MarginalRVs
+    fg, memo = fgraph_from_model(model)
+    unmarginalize_rewrite(fg)
+    unmarginalized_model = model_from_fgraph(fg, mutate_fgraph=True)
+    if rvs_to_unmarginalize is None:
+        return unmarginalized_model
+
+    # Re-marginalize the variables we want to keep marginalized
+    if not isinstance(rvs_to_unmarginalize, list | tuple):
+        rvs_to_unmarginalize = (rvs_to_unmarginalize,)
+    rvs_to_unmarginalize = set(rvs_to_unmarginalize)
+
+    old_free_rv_names = set(rv.name for rv in model.free_RVs)
+    new_free_rv_names = set(
+        rv.name for rv in unmarginalized_model.free_RVs if rv.name not in old_free_rv_names
+    )
+    if rvs_to_unmarginalize - new_free_rv_names:
+        raise ValueError(
+            f"Unrecognized rvs_to_unmarginalize: {rvs_to_unmarginalize - new_free_rv_names}"
+        )
+    rvs_to_keep_marginalized = tuple(new_free_rv_names - rvs_to_unmarginalize)
+    return marginalize(unmarginalized_model, rvs_to_keep_marginalized)
+
+
+def transform_posterior_pts(model, posterior_pts):
+    """Create a function from the untransformed space to the transformed space"""
+    # TODO: This should be a utility in PyMC
+    transformed_rvs = []
+    transformed_names = []
+
+    for rv in model.free_RVs:
+        transform = model.rvs_to_transforms.get(rv)
+        if transform is None:
+            transformed_rvs.append(rv)
+            transformed_names.append(rv.name)
         else:
-            seeds = [None] * len(vars_to_recover)
+            transformed_rv = transform.forward(rv, *rv.owner.inputs)
+            transformed_rvs.append(transformed_rv)
+            transformed_names.append(model.rvs_to_values[rv].name)
 
-        posterior = idata.posterior
+    fn = compile_pymc(
+        inputs=[In(inp, borrow=True) for inp in model.free_RVs],
+        outputs=[Out(out, borrow=True) for out in transformed_rvs],
+    )
+    fn.trust_input = True
 
-        # Remove Deterministics
-        posterior_values = posterior[
-            [rv.name for rv in self.free_RVs if rv not in self.marginalized_rvs]
-        ]
+    # TODO: This should work with vectorized inputs
+    return [dict(zip(transformed_names, fn(**point))) for point in posterior_pts]
 
-        sample_dims = ("chain", "draw")
-        posterior_pts, stacked_dims = dataset_to_point_list(posterior_values, sample_dims)
 
-        # Handle Transforms
-        transform_fn, transform_names = self._to_transformed()
+def recover_marginals(
+    model: Model,
+    idata: InferenceData,
+    var_names: Sequence[str] | None = None,
+    return_samples: bool = True,
+    extend_inferencedata: bool = True,
+    random_seed: RandomState = None,
+):
+    """Computes posterior log-probabilities and samples of marginalized variables
+    conditioned on parameters of the model given InferenceData with posterior group
 
-        def transform_input(inputs):
-            return dict(zip(transform_names, transform_fn(inputs)))
+    When there are multiple marginalized variables, each marginalized variable is
+    conditioned on both the parameters and the other variables still marginalized
 
-        posterior_pts = [transform_input(vs) for vs in posterior_pts]
+    All log-probabilities are within the transformed space
 
-        rv_dict = {}
-        rv_dims = {}
-        for seed, marginalized_rv in zip(seeds, vars_to_recover):
-            supported_dists = (Bernoulli, Categorical, DiscreteUniform)
-            if not isinstance(marginalized_rv.owner.op, supported_dists):
-                raise NotImplementedError(
-                    f"RV with distribution {marginalized_rv.owner.op} cannot be recovered. "
-                    f"Supported distribution include {supported_dists}"
-                )
+    Parameters
+    ----------
+    model: Model
+        PyMC model with marginalized variables to recover
+    idata : InferenceData
+        InferenceData with posterior group
+    var_names : sequence of str, optional
+        List of variable names for which to compute posterior log-probabilities and samples. Defaults to all marginalized variables
+    return_samples : bool, default True
+        If True, also return samples of the marginalized variables
+    extend_inferencedata : bool, default True
+        Whether to extend the original InferenceData or return a new one
+    random_seed: int, array-like of int or SeedSequence, optional
+        Seed used to generating samples
 
-            m = self.clone()
-            marginalized_rv = m.vars_to_clone[marginalized_rv]
-            m.unmarginalize([marginalized_rv])
-            dependent_rvs = find_conditional_dependent_rvs(marginalized_rv, m.basic_RVs)
-            logps = m.logp(vars=[marginalized_rv, *dependent_rvs], sum=False)
+    Returns
+    -------
+    idata : InferenceData
+        InferenceData with where a lp_{varname} and {varname} for each marginalized variable in var_names added to the posterior group
 
-            # Handle batch dims for marginalized value and its dependent RVs
-            dependent_rvs_dim_connections = subgraph_batch_dim_connection(
-                marginalized_rv, dependent_rvs
-            )
-            marginalized_logp, *dependent_logps = logps
-            joint_logp = marginalized_logp + reduce_batch_dependent_logps(
-                dependent_rvs_dim_connections,
-                [dependent_var.owner.op for dependent_var in dependent_rvs],
-                dependent_logps,
-            )
+    .. code-block:: python
 
-            marginalized_value = m.rvs_to_values[marginalized_rv]
-            other_values = [v for v in m.value_vars if v is not marginalized_value]
-
-            rv_shape = constant_fold(tuple(marginalized_rv.shape), raise_not_constant=False)
-            rv_domain = get_domain_of_finite_discrete_rv(marginalized_rv)
-            rv_domain_tensor = pt.moveaxis(
-                pt.full(
-                    (*rv_shape, len(rv_domain)),
-                    rv_domain,
-                    dtype=marginalized_rv.dtype,
-                ),
-                -1,
-                0,
-            )
+        import pymc as pm
+        from pymc_extras import MarginalModel
 
-            batched_joint_logp = vectorize_graph(
-                joint_logp,
-                replace={marginalized_value: rv_domain_tensor},
-            )
-            batched_joint_logp = pt.moveaxis(batched_joint_logp, 0, -1)
-
-            joint_logp_norm = log_softmax(batched_joint_logp, axis=-1)
-            if return_samples:
-                rv_draws = pymc.Categorical.dist(logit_p=batched_joint_logp)
-                if isinstance(marginalized_rv.owner.op, DiscreteUniform):
-                    rv_draws += rv_domain[0]
-                outputs = [joint_logp_norm, rv_draws]
-            else:
-                outputs = joint_logp_norm
-
-            rv_loglike_fn = compile_pymc(
-                inputs=other_values,
-                outputs=outputs,
-                on_unused_input="ignore",
-                random_seed=seed,
+        with MarginalModel() as m:
+            p = pm.Beta("p", 1, 1)
+            x = pm.Bernoulli("x", p=p, shape=(3,))
+            y = pm.Normal("y", pm.math.switch(x, -10, 10), observed=[10, 10, -10])
+
+            m.marginalize([x])
+
+            idata = pm.sample()
+            m.recover_marginals(idata, var_names=["x"])
+
+
+    """
+    unmarginal_model = unmarginalize(model)
+
+    # Find the names of the marginalized variables
+    model_var_names = set(rv.name for rv in model.free_RVs)
+    marginalized_rv_names = [
+        rv.name for rv in unmarginal_model.free_RVs if rv.name not in model_var_names
+    ]
+
+    if var_names is None:
+        var_names = marginalized_rv_names
+
+    var_names = [var if isinstance(var, str) else var.name for var in var_names]
+    var_names_to_recover = [name for name in marginalized_rv_names if name in var_names]
+    missing_names = [name for name in var_names_to_recover if name not in marginalized_rv_names]
+    if missing_names:
+        raise ValueError(f"Unrecognized var_names: {missing_names}")
+
+    if return_samples and random_seed is not None:
+        seeds = _get_seeds_per_chain(random_seed, len(var_names_to_recover))
+    else:
+        seeds = [None] * len(var_names_to_recover)
+
+    posterior_pts, stacked_dims = dataset_to_point_list(
+        # Remove Deterministics
+        idata.posterior[[rv.name for rv in model.free_RVs]],
+        sample_dims=("chain", "draw"),
+    )
+    transformed_posterior_pts = transform_posterior_pts(model, posterior_pts)
+
+    rv_dict = {}
+    rv_dims = {}
+    for seed, var_name_to_recover in zip(seeds, var_names_to_recover):
+        var_to_recover = unmarginal_model[var_name_to_recover]
+        supported_dists = (Bernoulli, Categorical, DiscreteUniform)
+        if not isinstance(var_to_recover.owner.op, supported_dists):
+            raise NotImplementedError(
+                f"RV with distribution {var_to_recover.owner.op} cannot be recovered. "
+                f"Supported distribution include {supported_dists}"
             )
 
-            logvs = [rv_loglike_fn(**vs) for vs in posterior_pts]
+        other_marginalized_rvs_names = marginalized_rv_names.copy()
+        other_marginalized_rvs_names.remove(var_name_to_recover)
+        dependent_rvs = [
+            rv
+            for rv in find_conditional_dependent_rvs(var_to_recover, unmarginal_model.basic_RVs)
+            if rv.name not in other_marginalized_rvs_names
+        ]
+        # Handle batch dims for marginalized value and its dependent RVs
+        dependent_rvs_dim_connections = subgraph_batch_dim_connection(var_to_recover, dependent_rvs)
+
+        marginalized_model = marginalize(unmarginal_model, other_marginalized_rvs_names)
 
-            if return_samples:
-                logps, samples = zip(*logvs)
-                logps = np.array(logps)
-                samples = np.array(samples)
-                rv_dict[marginalized_rv.name] = samples.reshape(
-                    tuple(len(coord) for coord in stacked_dims.values()) + samples.shape[1:],
-                )
-            else:
-                logps = np.array(logvs)
+        marginalized_var_to_recover = marginalized_model[var_name_to_recover]
+        dependent_rvs = [marginalized_model[rv.name] for rv in dependent_rvs]
 
-            rv_dict["lp_" + marginalized_rv.name] = logps.reshape(
-                tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:],
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=NonSeparableLogpWarning)
+            logps = marginalized_model.logp(
+                vars=[marginalized_var_to_recover, *dependent_rvs], sum=False
             )
-            if marginalized_rv.name in m.named_vars_to_dims:
-                rv_dims[marginalized_rv.name] = list(m.named_vars_to_dims[marginalized_rv.name])
-                rv_dims["lp_" + marginalized_rv.name] = rv_dims[marginalized_rv.name] + [
-                    "lp_" + marginalized_rv.name + "_dim"
-                ]
-
-        coords, dims = coords_and_dims_for_inferencedata(self)
-        dims.update(rv_dims)
-        rv_dataset = dict_to_dataset(
-            rv_dict,
-            library=pymc,
-            dims=dims,
-            coords=coords,
-            default_dims=list(sample_dims),
-            skip_event_dims=True,
-        )
 
-        if extend_inferencedata:
-            idata.posterior = idata.posterior.assign(rv_dataset)
-            return idata
-        else:
-            return rv_dataset
+        marginalized_logp, *dependent_logps = logps
+        joint_logp = marginalized_logp + reduce_batch_dependent_logps(
+            dependent_rvs_dim_connections,
+            [dependent_var.owner.op for dependent_var in dependent_rvs],
+            dependent_logps,
+        )
 
+        marginalized_value = marginalized_model.rvs_to_values[marginalized_var_to_recover]
+        other_values = [v for v in marginalized_model.value_vars if v is not marginalized_value]
+
+        rv_shape = constant_fold(tuple(var_to_recover.shape), raise_not_constant=False)
+        rv_domain = get_domain_of_finite_discrete_rv(var_to_recover)
+        rv_domain_tensor = pt.moveaxis(
+            pt.full(
+                (*rv_shape, len(rv_domain)),
+                rv_domain,
+                dtype=var_to_recover.dtype,
+            ),
+            -1,
+            0,
+        )
 
-def marginalize(model: Model, rvs_to_marginalize: ModelRVs) -> MarginalModel:
-    """Marginalize a subset of variables in a PyMC model.
+        batched_joint_logp = vectorize_graph(
+            joint_logp,
+            replace={marginalized_value: rv_domain_tensor},
+        )
+        batched_joint_logp = pt.moveaxis(batched_joint_logp, 0, -1)
+
+        joint_logp_norm = log_softmax(batched_joint_logp, axis=-1)
+        if return_samples:
+            rv_draws = Categorical.dist(logit_p=batched_joint_logp)
+            if isinstance(var_to_recover.owner.op, DiscreteUniform):
+                rv_draws += rv_domain[0]
+            outputs = [joint_logp_norm, rv_draws]
+        else:
+            outputs = joint_logp_norm
 
-    This creates a class of `MarginalModel` from an existing `Model`, with the specified
-    variables marginalized.
+        rv_loglike_fn = compile_pymc(
+            inputs=other_values,
+            outputs=outputs,
+            on_unused_input="ignore",
+            random_seed=seed,
+        )
 
-    See documentation for `MarginalModel` for more information.
+        logvs = [rv_loglike_fn(**vs) for vs in transformed_posterior_pts]
 
-    Parameters
-    ----------
-    model : Model
-        PyMC model to marginalize. Original variables well be cloned.
-    rvs_to_marginalize : Sequence[TensorVariable]
-        Variables to marginalize in the returned model.
+        if return_samples:
+            logps, samples = zip(*logvs)
+            logps = np.asarray(logps)
+            samples = np.asarray(samples)
+            rv_dict[var_name_to_recover] = samples.reshape(
+                tuple(len(coord) for coord in stacked_dims.values()) + samples.shape[1:],
+            )
+        else:
+            logps = np.asarray(logvs)
 
-    Returns
-    -------
-    marginal_model: MarginalModel
-        Marginal model with the specified variables marginalized.
-    """
-    if not isinstance(rvs_to_marginalize, tuple | list):
-        rvs_to_marginalize = (rvs_to_marginalize,)
-    rvs_to_marginalize = [rv if isinstance(rv, str) else rv.name for rv in rvs_to_marginalize]
+        rv_dict["lp_" + var_name_to_recover] = logps.reshape(
+            tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:],
+        )
+        if var_name_to_recover in unmarginal_model.named_vars_to_dims:
+            rv_dims[var_name_to_recover] = list(
+                unmarginal_model.named_vars_to_dims[var_name_to_recover]
+            )
+            rv_dims["lp_" + var_name_to_recover] = rv_dims[var_name_to_recover] + [
+                "lp_" + var_name_to_recover + "_dim"
+            ]
+
+    coords, dims = coords_and_dims_for_inferencedata(unmarginal_model)
+    dims.update(rv_dims)
+    rv_dataset = dict_to_dataset(
+        rv_dict,
+        library=pymc,
+        dims=dims,
+        coords=coords,
+        skip_event_dims=True,
+    )
 
-    marginal_model = MarginalModel.from_model(model)
-    marginal_model.marginalize(rvs_to_marginalize)
-    return marginal_model
+    if extend_inferencedata:
+        idata.posterior = idata.posterior.assign(rv_dataset)
+        return idata
+    else:
+        return rv_dataset
 
 
 def collect_shared_vars(outputs, blockers):
     return [
-        inp for inp in graph_inputs(outputs, blockers=blockers) if isinstance(inp, SharedVariable)
+        inp
+        for inp in graph_inputs(outputs, blockers=blockers)
+        if (isinstance(inp, SharedVariable) and inp not in blockers)
     ]
 
 
-def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs):
-    dependent_rvs = find_conditional_dependent_rvs(rv_to_marginalize, all_rvs)
-    if not dependent_rvs:
-        raise ValueError(f"No RVs depend on marginalized RV {rv_to_marginalize}")
+def remove_model_vars(vars):
+    """Remove ModelVars from the graph of vars."""
+    model_vars = [var for var in vars if isinstance(var.owner.op, ModelValuedVar)]
+    replacements = [(model_var, model_var.owner.inputs[0]) for model_var in model_vars]
+    fgraph = FunctionGraph(outputs=vars, clone=False)
+    toposort_replace(fgraph, replacements)
+    return fgraph.outputs
 
-    marginalized_rv_input_rvs = find_conditional_input_rvs([rv_to_marginalize], all_rvs)
-    other_direct_rv_ancestors = [
-        rv
-        for rv in find_conditional_input_rvs(dependent_rvs, all_rvs)
-        if rv is not rv_to_marginalize
-    ]
 
+def replace_finite_discrete_marginal_subgraph(
+    fgraph, rv_to_marginalize, dependent_rvs, input_rvs
+) -> None:
     # If the marginalized RV has multiple dimensions, check that graph between
     # marginalized RV and dependent RVs does not mix information from batch dimensions
     # (otherwise logp would require enumerating over all combinations of batch dimension values)
@@ -574,22 +567,42 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
             "You can try splitting the marginalized RV into separate components and marginalizing them separately."
         ) from e
 
-    input_rvs = list(set((*marginalized_rv_input_rvs, *other_direct_rv_ancestors)))
     output_rvs = [rv_to_marginalize, *dependent_rvs]
+    rng_updates = collect_default_updates(output_rvs, inputs=input_rvs, must_be_shared=False)
+    outputs = output_rvs + list(rng_updates.values())
+    inputs = input_rvs + list(rng_updates.keys())
+    # Add any other shared variable inputs
+    inputs += collect_shared_vars(output_rvs, blockers=inputs)
 
-    # We are strict about shared variables in SymbolicRandomVariables
-    inputs = input_rvs + collect_shared_vars(output_rvs, blockers=input_rvs)
+    inner_inputs = [inp.clone() for inp in inputs]
+    inner_outputs = clone_replace(outputs, replace=dict(zip(inputs, inner_inputs)))
+    inner_outputs = remove_model_vars(inner_outputs)
 
-    if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
+    if isinstance(inner_outputs[0].owner.op, DiscreteMarkovChain):
         marginalize_constructor = MarginalDiscreteMarkovChainRV
     else:
         marginalize_constructor = MarginalFiniteDiscreteRV
 
+    _, _, *dims = rv_to_marginalize.owner.inputs
     marginalization_op = marginalize_constructor(
-        inputs=inputs,
-        outputs=output_rvs,  # TODO: Add RNG updates to outputs so this can be used in the generative graph
+        inputs=inner_inputs,
+        outputs=inner_outputs,
         dims_connections=dependent_rvs_dim_connections,
+        dims=dims,
     )
-    new_output_rvs = marginalization_op(*inputs)
-    fgraph.replace_all(tuple(zip(output_rvs, new_output_rvs)))
-    return output_rvs, new_output_rvs
+
+    new_outputs = marginalization_op(*inputs)
+    for old_output, new_output in zip(outputs, new_outputs):
+        new_output.name = old_output.name
+
+    model_replacements = []
+    for old_output, new_output in zip(outputs, new_outputs):
+        if old_output is rv_to_marginalize or not isinstance(old_output.owner.op, ModelValuedVar):
+            # Replace the marginalized ModelFreeRV (or non model-variables) themselves
+            var_to_replace = old_output
+        else:
+            # Replace the underlying RV, keeping the same value, transform and dims
+            var_to_replace = old_output.owner.inputs[0]
+        model_replacements.append((var_to_replace, new_output))
+
+    fgraph.replace_all(model_replacements)