Repository: pyro-ppl/pyro

Issue:
Predictive class fails for non-scalar sites
I usually use raw poutines to for prediction:
```py
vectorize = pyro.plate("_vectorize", num_samples, dim=-2)
trace = poutine.trace(vectorize(guide)).get_trace()
samples = poutine.replay(vectorize(model), trace)()
```
When instead I try the new `Predictive` class, this fails on models with nontrivial plate structure:
```py
predictive = Predictive(model, guide=guide, return_sites=["x", "y"],
                        num_samples=num_samples, parallel=True)
samples = predictive.get_samples()
```
I believe this is due to reshaping of posterior samples without regard to their event shape. I don't know why `_predictive()` even needs to reshape samples.

I have added a failing test in this branch. Feel free to take over the branch and create a PR:
https://github.com/pyro-ppl/pyro/pull/new/fix-preditive-shapes


Discussion:

Resolving patch:
diff --git a/pyro/infer/predictive.py b/pyro/infer/predictive.py
--- a/pyro/infer/predictive.py
+++ b/pyro/infer/predictive.py
@@ -28,7 +28,7 @@ def _guess_max_plate_nesting(model, args, kwargs):
 
 
 def _predictive_sequential(model, posterior_samples, model_args, model_kwargs,
-                           num_samples, sample_sites, return_trace=False):
+                           num_samples, return_site_shapes, 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):
@@ -36,10 +36,13 @@ def _predictive_sequential(model, posterior_samples, model_args, model_kwargs,
         if return_trace:
             collected.append(trace)
         else:
-            collected.append({site: trace.nodes[site]['value'] for site in sample_sites})
+            collected.append({site: trace.nodes[site]['value'] for site in return_site_shapes})
 
-    return collected if return_trace else {site: torch.stack([s[site] for s in collected])
-                                           for site in sample_sites}
+    if return_trace:
+        return collected
+    else:
+        return {site: torch.stack([s[site] for s in collected]).reshape(shape)
+                for site, shape in return_site_shapes.items()}
 
 
 def _predictive(model, posterior_samples, num_samples, return_sites=None,
@@ -61,7 +64,8 @@ def _predictive(model, posterior_samples, num_samples, return_sites=None,
 
     return_site_shapes = {}
     for site in model_trace.stochastic_nodes + model_trace.observation_nodes:
-        site_shape = (num_samples,) + model_trace.nodes[site]['value'].shape
+        append_ndim = max_plate_nesting - len(model_trace.nodes[site]["fn"].batch_shape)
+        site_shape = (num_samples,) + (1,) * append_ndim + model_trace.nodes[site]['value'].shape
         if isinstance(return_sites, (list, tuple, set)):
             if site in return_sites:
                 return_site_shapes[site] = site_shape
@@ -77,7 +81,7 @@ def _predictive(model, posterior_samples, num_samples, return_sites=None,
 
     if not parallel:
         return _predictive_sequential(model, posterior_samples, model_args, model_kwargs, num_samples,
-                                      return_site_shapes.keys(), return_trace=False)
+                                      return_site_shapes, return_trace=False)
 
     trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\
         .get_trace(*model_args, **model_kwargs)


Test patch:
diff --git a/tests/infer/test_predictive.py b/tests/infer/test_predictive.py
--- a/tests/infer/test_predictive.py
+++ b/tests/infer/test_predictive.py
@@ -92,3 +92,30 @@ def test_posterior_predictive_svi_one_hot():
     posterior_predictive = Predictive(one_hot_model, posterior_samples)
     marginal_return_vals = posterior_predictive.get_samples(pseudocounts)["obs"]
     assert_close(marginal_return_vals.mean(dim=0), true_probs.unsqueeze(0), rtol=0.1)
+
+
+@pytest.mark.parametrize("parallel", [False, True])
+def test_shapes(parallel):
+    num_samples = 10
+
+    def model():
+        x = pyro.sample("x", dist.Normal(0, 1).expand([2]).to_event(1))
+        with pyro.plate("plate", 5):
+            loc, log_scale = x.unbind(-1)
+            y = pyro.sample("y", dist.Normal(loc, log_scale.exp()))
+        return dict(x=x, y=y)
+
+    guide = AutoDiagonalNormal(model)
+
+    # Compute by hand.
+    vectorize = pyro.plate("_vectorize", num_samples, dim=-2)
+    trace = poutine.trace(vectorize(guide)).get_trace()
+    expected = poutine.replay(vectorize(model), trace)()
+
+    # Use Predictive.
+    predictive = Predictive(model, guide=guide, return_sites=["x", "y"],
+                            num_samples=num_samples, parallel=parallel)
+    actual = predictive.get_samples()
+    assert set(actual) == set(expected)
+    assert actual["x"].shape == expected["x"].shape
+    assert actual["y"].shape == expected["y"].shape