Repository: pyro-ppl/pyro

Issue:
Uncollapse compound distributions in traces generated from MCMC
**NOTE:** This is a sketch, and we will need to iron out the API issues before this can be merged.

This adds the ability to collapse latent sites in compound distributions (i.e. latent observed conjugate pair) during inference in HMC, and uncollapse it later by sampling from the exact posterior distribution followed by an observed site with the sampled parameters. Note that this requires the user to use compound distributions directly in the model. Later, when we can do conjugacy detection via `config_infer`, this can be relaxed.

**Design**: The original implementation could not handle latents within arbitrary plates, which has been rectified now. This has however added some complexity to these classed due to which I have moved them to `contrib`.  
 - We are passing in the name of the site to the constructor for the distribution, and have two methods, `pin_latent` and `observe` to specify the location for the collapsed latent site and the observed site. e.g. below. Note that the reason to pass in the name is to be able to do site matching later when sampling from the posterior, but an alternate design might be for the user to specify the latent and observed sample in their models with the requirement that they use a pre-established naming convention. Another complication with this design choice would be dynamically adjusting the shape of the latent (which is currently happening in `.pin_latent()`). EDIT: I think the latter is a much better design choice though, and I am exploring work-arounds to not have to pass in name into the constructor.

   ```python
   b = BetaBinomial("betabinom", ...)
   with pyro.plate("dim0", ...):
       b.pin_latent()  # Adds a collapsed site named "betabinom.latent"
       ...
       with pyro.plate("data", ...):
           b.observe(data)  # Add a sample(observed) site called "betabinom"
   ```
 - When we call `.pin_latent()`, the latent size is adjusted so that all the other methods like `.sample` and `.log_prob` use the correctly shaped latent distribution.

**Changes:**
 - Introduces `_posterior_latent` and `_posterior_predictive` methods to conjugate distributions - the former is the posterior distribution of the parameter values given the observation, and the latter takes in parameters sampled from the posterior to give the posterior predictive distribution. Note that this uncollapsing will be done by the user outside of HMC, to retrieve samples from the  posterior distribution of the (originally collapsed) latent site. 
 - Minor changes to the HMC class so that in case the model only comprises of compound distributions (i.e. there are only observe statements), we do not fail but return a deterministic trace.
 - An `infer_conjugate` poutine that collapses latent sites with the infer kwarg `collapse=True`, , and uncollapses these by sampling them from the exact posterior when the model is wrapped under `infer_conjugate` poutine.


Discussion:
~Please hold off on reviewing this - after discussion with @fritzo, we concluded that we will need more information in our compound distributions to infer the shape of the latent, and this implementation isn't general in that it assumes that the latent shape is the shape of the original parameters that was used to construct the distribution.~ 

This has been rectified, so opening this again for discussions on improving the interface.
@fritzo - This has been updated with the changes that I mentioned that allow for sampling latent values inside arbitrary plates, but this has forced me to make certain design decisions that don't align well with the rest of distributions. I plan to use this branch for the MOBB model, but would like to discuss if you see a way we can integrate this into `distributions/conjugate.py`.
I will put up a new PR with a simplified interface for discussion purpose. 

Resolving patch:
diff --git a/pyro/contrib/conjugate/__init__.py b/pyro/contrib/conjugate/__init__.py
new file mode 100644
diff --git a/pyro/contrib/conjugate/infer.py b/pyro/contrib/conjugate/infer.py
new file mode 100644
--- /dev/null
+++ b/pyro/contrib/conjugate/infer.py
@@ -0,0 +1,164 @@
+import pyro.distributions as dist
+from pyro.distributions.util import sum_leftmost
+from pyro.poutine.messenger import Messenger
+from pyro.poutine.replay_messenger import ReplayMessenger
+
+
+def _make_cls(base, static_attrs, instance_attrs, parent_linkage=None):
+    r"""
+    Dynamically create classes named `_ + base.__name__`, which extend the
+    base class with other optional instance and class attributes, and have
+    a custom `.expand` method to propagate these attributes on expanded
+    instances.
+
+    :param cls base: Base class.
+    :param dict static_attrs: static attributes to add to class.
+    :param dict instance_attrs: instance attributes for initialization.
+    :param str parent_linkage: attribute in the parent class that holds
+        a reference to the distribution class.
+    :return cls: dynamically generated class.
+    """
+    def _expand(self, batch_shape, _instance=None):
+        new = self._get_checked_instance(cls, _instance)
+        for attr in instance_attrs:
+            setattr(new, attr, getattr(self, attr))
+        if parent_linkage:
+            setattr(new.parent, parent_linkage, new)
+        return base.expand(self, batch_shape, _instance=new)
+
+    name = "_" + base.__name__
+    cls = type(name, (base,), instance_attrs)
+    for k, v in static_attrs.items():
+        setattr(cls, k, v)
+    cls.expand = _expand
+    return cls
+
+
+def _latent(base, parent):
+    return _make_cls(base, {"collapsible": True}, {"site_name": None, "parent": parent}, "_latent")
+
+
+def _conditional(base, parent):
+    return _make_cls(base, {"marginalize_latent": True}, {"parent": parent}, "_conditional")
+
+
+def _compound(base, parent):
+    return _make_cls(base, {}, {"parent": parent})
+
+
+class BetaBinomialPair(object):
+    def __init__(self):
+        self._latent = None
+        self._conditional = None
+
+    def latent(self, *args, **kwargs):
+        self._latent = _latent(dist.Beta, parent=self)(*args, **kwargs)
+        return self._latent
+
+    def conditional(self, *args, **kwargs):
+        self._conditional = _conditional(dist.Binomial, parent=self)(*args, **kwargs)
+        return self._conditional
+
+    def posterior(self, obs):
+        concentration1 = self._latent.concentration1
+        concentration0 = self._latent.concentration0
+        total_count = self._conditional.total_count
+        reduce_dims = len(obs.size()) - len(concentration1.size())
+        # Unexpand total_count to have the same shape as concentration0.
+        # Raise exception if this isn't possible.
+        total_count = sum_leftmost(total_count, reduce_dims)
+        summed_obs = sum_leftmost(obs, reduce_dims)
+        return dist.Beta(concentration1 + summed_obs,
+                         total_count + concentration0 - summed_obs,
+                         validate_args=self._latent._validate_args)
+
+    def compound(self):
+        return _compound(dist.BetaBinomial, parent=self)(concentration1=self._latent.concentration1,
+                                                         concentration0=self._latent.concentration0,
+                                                         total_count=self._conditional.total_count)
+
+
+class GammaPoissonPair(object):
+    def __init__(self):
+        self._latent = None
+        self._conditional = None
+
+    def latent(self, *args, **kwargs):
+        self._latent = _latent(dist.Gamma, parent=self)(*args, **kwargs)
+        return self._latent
+
+    def conditional(self, *args, **kwargs):
+        self._conditional = _conditional(dist.Poisson, parent=self)(*args, **kwargs)
+        return self._conditional
+
+    def posterior(self, obs):
+        concentration = self._latent.concentration
+        rate = self._latent.rate
+        reduce_dims = len(obs.size()) - len(rate.size())
+        num_obs = obs.shape[:reduce_dims].numel()
+        summed_obs = sum_leftmost(obs, reduce_dims)
+        return dist.Gamma(concentration + summed_obs, rate + num_obs)
+
+    def compound(self):
+        return _compound(dist.GammaPoisson, parent=self)(concentration=self._latent.concentration,
+                                                         rate=self._latent.rate)
+
+
+class UncollapseConjugateMessenger(ReplayMessenger):
+    r"""
+    Extends `~pyro.poutine.replay_messenger.ReplayMessenger` to uncollapse
+    compound distributions.
+    """
+    def _pyro_sample(self, msg):
+        is_collapsible = getattr(msg["fn"], "collapsible", False)
+        if is_collapsible:
+            conj_node, parent = None, None
+            for site_name in self.trace.observation_nodes + self.trace.stochastic_nodes:
+                parent = getattr(self.trace.nodes[site_name]["fn"], "parent")
+                if parent is not None and parent._latent.site_name == msg["name"]:
+                    conj_node = self.trace.nodes[site_name]
+                    break
+            assert conj_node is not None, "Collapsible latent site `{}` with no corresponding conjugate site."\
+                .format(msg["name"])
+            msg["fn"] = parent.posterior(conj_node["value"])
+            msg["value"] = msg["fn"].sample()
+        else:
+            return super(UncollapseConjugateMessenger, self)._pyro_sample(msg)
+
+
+def uncollapse_conjugate(fn=None, trace=None, params=None):
+    r"""
+    This extends the behavior of :function:`~pyro.poutine.replay` poutine, so that in
+    addition to replaying the values at sample sites from the ``trace`` in the
+    original callable ``fn`` when the same sites are sampled, this also "uncollapses"
+    any observed compound distributions (defined in :module:`pyro.distributions.conjugate`)
+    by sampling the originally collapsed parameter values from its posterior distribution
+    followed by observing the data with the sampled parameter values.
+    """
+    msngr = UncollapseConjugateMessenger(trace, params)
+    return msngr(fn) if fn is not None else msngr
+
+
+class CollapseConjugateMessenger(Messenger):
+    def _pyro_sample(self, msg):
+        is_collapsible = getattr(msg["fn"], "collapsible", False)
+        marginalize_latent = getattr(msg["fn"], "marginalize_latent", False)
+        if is_collapsible:
+            msg["fn"].site_name = msg["name"]
+            msg["stop"] = True
+        elif marginalize_latent:
+            msg["fn"] = msg["fn"].parent.compound()
+        else:
+            return
+
+
+def collapse_conjugate(fn=None):
+    r"""
+    This replaces a latent-observed pair by collapsing the latent site
+    (whose distribution has attribute `collapsible=True`), and replacing the
+    observed site (whose distribution has attribute `marginalize_latent=True`)
+    with a compound probability distribution that marginalizes out the latent
+    site.
+    """
+    msngr = CollapseConjugateMessenger()
+    return msngr(fn) if fn is not None else msngr


Test patch:
diff --git a/tests/infer/mcmc/test_nuts.py b/tests/infer/mcmc/test_nuts.py
--- a/tests/infer/mcmc/test_nuts.py
+++ b/tests/infer/mcmc/test_nuts.py
@@ -10,6 +10,7 @@
 import pyro
 import pyro.distributions as dist
 from pyro.contrib.autoguide import AutoDelta
+from pyro.contrib.conjugate.infer import BetaBinomialPair, collapse_conjugate, uncollapse_conjugate, GammaPoissonPair
 from pyro.infer import TraceEnum_ELBO, SVI
 from pyro.infer.mcmc.mcmc import MCMC
 from pyro.infer.mcmc.nuts import NUTS
@@ -318,3 +319,47 @@ def _generate_data():
     if num_steps == 30:
         nuts_kernel.initial_trace = _get_initial_trace()
     MCMC(nuts_kernel, num_samples=5, warmup_steps=5).run(data)
+
+
+@pytest.mark.parametrize("hyperpriors", [False, True])
+def test_beta_binomial(hyperpriors):
+    def model(data):
+        with pyro.plate("plate_0", data.shape[-1]):
+            alpha = pyro.sample("alpha", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor([1., 1.])
+            beta = pyro.sample("beta", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor([1., 1.])
+            beta_binom = BetaBinomialPair()
+            with pyro.plate("plate_1", data.shape[-2]):
+                probs = pyro.sample("probs", beta_binom.latent(alpha, beta))
+                with pyro.plate("data", data.shape[0]):
+                    pyro.sample("binomial", beta_binom.conditional(probs=probs, total_count=total_count), obs=data)
+
+    true_probs = torch.tensor([[0.7, 0.4], [0.6, 0.4]])
+    total_count = torch.tensor([[1000, 600], [400, 800]])
+    data = dist.Binomial(total_count=total_count, probs=true_probs).sample(sample_shape=(torch.Size((10,))))
+    hmc_kernel = NUTS(collapse_conjugate(model), jit_compile=True, ignore_jit_warnings=True)
+    mcmc_run = MCMC(hmc_kernel, num_samples=80, warmup_steps=50).run(data)
+    mcmc_run.exec_traces = [poutine.trace(uncollapse_conjugate(model, tr)).get_trace(data)
+                            for tr in mcmc_run.exec_traces]
+    posterior = mcmc_run.marginal(["probs"]).empirical["probs"]
+    assert_equal(posterior.mean, true_probs, prec=0.05)
+
+
+@pytest.mark.parametrize("hyperpriors", [False, True])
+def test_gamma_poisson(hyperpriors):
+    def model(data):
+        with pyro.plate("latent_dim", data.shape[1]):
+            alpha = pyro.sample("alpha", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor([1., 1.])
+            beta = pyro.sample("beta", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor([1., 1.])
+            gamma_poisson = GammaPoissonPair()
+            rate = pyro.sample("rate", gamma_poisson.latent(alpha, beta))
+            with pyro.plate("data", data.shape[0]):
+                pyro.sample("obs", gamma_poisson.conditional(rate), obs=data)
+
+    true_rate = torch.tensor([3., 10.])
+    data = dist.Poisson(rate=true_rate).sample(sample_shape=(torch.Size((100,))))
+    hmc_kernel = NUTS(collapse_conjugate(model), jit_compile=True, ignore_jit_warnings=True)
+    mcmc_run = MCMC(hmc_kernel, num_samples=100, warmup_steps=50).run(data)
+    mcmc_run.exec_traces = [poutine.trace(uncollapse_conjugate(model, tr)).get_trace(data)
+                            for tr in mcmc_run.exec_traces]
+    posterior = mcmc_run.marginal(["rate"]).empirical["rate"]
+    assert_equal(posterior.mean, true_rate, prec=0.2)