Repository: pyro-ppl/pyro

Issue:
Add option for sequential prediction to `mcmc.predictive`
Based on this [forum](https://forum.pyro.ai/t/predictive-sampling-with-logistic-regression/1189/5) post, I think a lot of new users of MCMC struggle with using plate handlers to correctly annotate their batch dimensions (which isn't strictly needed for doing inference). 

While this isn't necessarily a requirement for MCMC, it becomes important if users would like to use the `predictive` utility to do predictions after running MCMC. This utility wraps the model in an outermost plate to vectorize predictions. I think we should support a `parallel=False` argument to this utility which by default will run the model forward on a per trace basis and only use `plate` if parallel is set to True. This will make it easier for new users to play around with MCMC who may not be well acquainted with our auto batching semantics.





Discussion:

Resolving patch:
diff --git a/examples/baseball.py b/examples/baseball.py
--- a/examples/baseball.py
+++ b/examples/baseball.py
@@ -228,7 +228,7 @@ def evaluate_log_posterior_density(model, posterior_samples, baseball_dataset):
     at_bats_season, hits_season = test[:, 0], test[:, 1]
     with ignore_experimental_warning():
         trace = predictive(model, posterior_samples, at_bats_season, hits_season,
-                           return_trace=True)
+                           parallel=True, return_trace=True)
     # Use LogSumExp trick to evaluate $log(1/num_samples \sum_i p(new_data | \theta^{i})) $,
     # where $\theta^{i}$ are parameter samples from the model's posterior.
     trace.compute_log_prob()
diff --git a/pyro/infer/mcmc/util.py b/pyro/infer/mcmc/util.py
--- a/pyro/infer/mcmc/util.py
+++ b/pyro/infer/mcmc/util.py
@@ -469,6 +469,21 @@ def summary(samples, prob=0.9, num_chains=1):
     print('\n')
 
 
+def _predictive_sequential(model, posterior_samples, model_args, model_kwargs,
+                           num_samples, sample_sites, return_trace=False):
+    collected = []
+    samples = [{k: v[i] for k, v in posterior_samples.items()} for i in range(num_samples)]
+    for i in range(num_samples):
+        trace = poutine.trace(poutine.condition(model, samples[i])).get_trace(*model_args, **model_kwargs)
+        if return_trace:
+            collected.append(trace)
+        else:
+            collected.append({site: trace.nodes[site]['value'] for site in sample_sites})
+
+    return collected if return_trace else {site: torch.stack([s[site] for s in collected])
+                                           for site in sample_sites}
+
+
 def predictive(model, posterior_samples, *args, **kwargs):
     """
     Run model by sampling latent parameters from `posterior_samples`, and return
@@ -494,6 +509,9 @@ def predictive(model, posterior_samples, *args, **kwargs):
           in `posterior_samples` are returned.
         * **return_trace** (``bool``) - whether to return the full trace. Note that this is vectorized
           over `num_samples`.
+        * **parallel** (``bool``) - predict in parallel by wrapping the existing model
+          in an outermost `plate` messenger. Note that this requires that the model has
+          all batch dims correctly annotated via :class:`~pyro.plate`. Default is `False`.
 
     :return: dict of samples from the predictive distribution, or a single vectorized
         `trace` (if `return_trace=True`).
@@ -503,6 +521,7 @@ def predictive(model, posterior_samples, *args, **kwargs):
     num_samples = kwargs.pop('num_samples', None)
     return_sites = kwargs.pop('return_sites', None)
     return_trace = kwargs.pop('return_trace', False)
+    parallel = kwargs.pop('parallel', False)
 
     max_plate_nesting = _guess_max_plate_nesting(model, args, kwargs)
     model_trace = prune_subsample_sites(poutine.trace(model).get_trace(*args, **kwargs))
@@ -537,6 +556,10 @@ def predictive(model, posterior_samples, *args, **kwargs):
             if site not in reshaped_samples:
                 return_site_shapes[site] = site_shape
 
+    if not parallel:
+        return _predictive_sequential(model, posterior_samples, args, kwargs, num_samples,
+                                      return_site_shapes.keys(), return_trace)
+
     def _vectorized_fn(fn):
         """
         Wraps a callable inside an outermost :class:`~pyro.plate` to parallelize


Test patch:
diff --git a/tests/infer/mcmc/test_mcmc_util.py b/tests/infer/mcmc/test_mcmc_util.py
--- a/tests/infer/mcmc/test_mcmc_util.py
+++ b/tests/infer/mcmc/test_mcmc_util.py
@@ -25,7 +25,8 @@ def model(data=None):
 
 
 @pytest.mark.parametrize("num_samples", [100, 200, None])
-def test_predictive(num_samples):
+@pytest.mark.parametrize("parallel", [False, True])
+def test_predictive(num_samples, parallel):
     model, data, true_probs = beta_bernoulli()
     init_params, potential_fn, transforms, _ = initialize_model(model,
                                                                 model_args=(data,))
@@ -40,7 +41,8 @@ def test_predictive(num_samples):
         with optional(pytest.warns(UserWarning), num_samples not in (None, 100)):
             predictive_samples = predictive(model, samples,
                                             num_samples=num_samples,
-                                            return_sites=["beta", "obs"])
+                                            return_sites=["beta", "obs"],
+                                            parallel=parallel)
 
     # check shapes
     assert predictive_samples["beta"].shape == (100, 5)