To add a new kind of marginalization (e.g., a new conjugate pair or approximation), you must implement a new module under pymc_extras/model/marginal/distributions/. You need to follow these four steps:
1. Subclass MarginalRV
Create a class inheriting from MarginalRV. The inner graph outputs must be laid out as [marginalized_rv, *dependent_rvs, *rng_updates].
2. Write the Rewrite
Create a function decorated with @node_rewriter(tracks=[MarginalSubgraph]). This function should:
- Use
extract_marginal_subgraph(node) to get inputs and outputs. - Inspect the nodes to see if they match your specific distribution pattern.
- Return
None if the pattern does not match (to allow other rewrites to run). - If it matches, return the new
MarginalRV outputs. - Register the rewrite using
marginal_rewrites_db.register(name, rewrite_func, priority).
3. Register the Marginal Log-Probability
Use the @_logprob.register(YourMarginalRVClass) decorator. Inside, use inline_ofg_outputs(op, inputs) to recover the inner generative graph and calculate the log-probability of the dependent values with the marginalized variable integrated out.
4. (Optional) Register the Conditional
To support conditional() and recover(), use the @marginalized_conditional.register(YourMarginalRVClass) decorator. This function receives the node's inputs and the dep_rvs (the variables the dependents are conditioned on) and must return a random variable representing $p( ext{marginalized} ext{ | } ext{dependents})$.
# 1. Subclass MarginalRV
from pymc_extras.model.marginal.distributions.core import MarginalRV
class GammaPoissonMarginalRV(MarginalRV):
def __init__(self, *args, marginalized_dims, **kwargs):
self.marginalized_dims = marginalized_dims
self.n_dependent_rvs = 1
super().__init__(*args, **kwargs)
# 2. Write the rewrite
from pymc.distributions import Gamma, Poisson
from pytensor.graph import node_rewriter
from pymc_extras.model.marginal.rewrites import (
MarginalSubgraph,
extract_marginal_subgraph,
marginal_rewrites_db,
)
@node_rewriter(tracks=[MarginalSubgraph])
def gamma_poisson_marginal(fgraph, node):
if node.op.n_dependent_rvs != 1:
return None
inputs, outputs = extract_marginal_subgraph(node)
marginalized_rv, dependent_rv = outputs[:2]
if not (
isinstance(marginalized_rv.owner.op, Gamma)
and isinstance(dependent_rv.owner.op, Poisson)
and marginalized_rv.type.ndim == 0
):
return None
[poisson_mu] = dependent_rv.owner.op.dist_params(dependent_rv.owner)
if poisson_mu is not marginalized_rv:
return None
typed_op = GammaPoissonMarginalRV(
inputs=inputs,
outputs=outputs,
marginalized_dims=node.op.marginalized_dims,
)
new_outputs = typed_op(*inputs)
return list(new_outputs)[: len(node.outputs)]
marginal_rewrites_db.register("gamma_poisson_marginal", gamma_poisson_marginal, "basic")
# 3. Register the marginal logp
from pymc import NegativeBinomial
from pymc.logprob.abstract import _logprob
from pymc.logprob.basic import logp
from pymc_extras.model.marginal.distributions.core import inline_ofg_outputs
@_logprob.register(GammaPoissonMarginalRV)
def gamma_poisson_marginal_logp(op, values, *inputs, **kwargs):
[value] = values
marginalized_rv, _ = inline_ofg_outputs(op, inputs)[:2]
alpha, scale = marginalized_rv.owner.op.dist_params(marginalized_rv.owner)
beta = 1 / scale
return logp(NegativeBinomial.dist(n=alpha, p=beta / (beta + 1)), value)
# 4. (Optional) Register the conditional
from pymc_extras.model.marginal.distributions.core import marginalized_conditional
@marginalized_conditional.register(GammaPoissonMarginalRV)
def gamma_poisson_conditional(op, inputs, dep_rvs):
marginalized = inline_ofg_outputs(op, inputs)[0]
alpha, scale = marginalized.owner.op.dist_params(marginalized.owner)
beta = 1 / scale
[dep_rv] = dep_rvs
return Gamma.dist(alpha + dep_rv, beta + 1)
