We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/infer/predictive.py`
Content:
```
1 from functools import reduce
2 import warnings
3 
4 import torch
5 
6 import pyro
7 import pyro.poutine as poutine
8 from pyro.poutine.util import prune_subsample_sites
9 
10 
11 def _guess_max_plate_nesting(model, args, kwargs):
12     """
13     Guesses max_plate_nesting by running the model once
14     without enumeration. This optimistically assumes static model
15     structure.
16     """
17     with poutine.block():
18         model_trace = poutine.trace(model).get_trace(*args, **kwargs)
19     sites = [site for site in model_trace.nodes.values()
20              if site["type"] == "sample"]
21 
22     dims = [frame.dim
23             for site in sites
24             for frame in site["cond_indep_stack"]
25             if frame.vectorized]
26     max_plate_nesting = -min(dims) if dims else 0
27     return max_plate_nesting
28 
29 
30 def _predictive_sequential(model, posterior_samples, model_args, model_kwargs,
31                            num_samples, sample_sites, return_trace=False):
32     collected = []
33     samples = [{k: v[i] for k, v in posterior_samples.items()} for i in range(num_samples)]
34     for i in range(num_samples):
35         trace = poutine.trace(poutine.condition(model, samples[i])).get_trace(*model_args, **model_kwargs)
36         if return_trace:
37             collected.append(trace)
38         else:
39             collected.append({site: trace.nodes[site]['value'] for site in sample_sites})
40 
41     return collected if return_trace else {site: torch.stack([s[site] for s in collected])
42                                            for site in sample_sites}
43 
44 
45 def _predictive(model, posterior_samples, num_samples, return_sites=None,
46                 return_trace=False, parallel=False, model_args=(), model_kwargs={}):
47     max_plate_nesting = _guess_max_plate_nesting(model, model_args, model_kwargs)
48     vectorize = pyro.plate("_num_predictive_samples", num_samples, dim=-max_plate_nesting-1)
49     model_trace = prune_subsample_sites(poutine.trace(model).get_trace(*model_args, **model_kwargs))
50     reshaped_samples = {}
51 
52     for name, sample in posterior_samples.items():
53         sample_shape = sample.shape[1:]
54         sample = sample.reshape((num_samples,) + (1,) * (max_plate_nesting - len(sample_shape)) + sample_shape)
55         reshaped_samples[name] = sample
56 
57     if return_trace:
58         trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\
59             .get_trace(*model_args, **model_kwargs)
60         return trace
61 
62     return_site_shapes = {}
63     for site in model_trace.stochastic_nodes + model_trace.observation_nodes:
64         site_shape = (num_samples,) + model_trace.nodes[site]['value'].shape
65         if isinstance(return_sites, (list, tuple, set)):
66             if site in return_sites:
67                 return_site_shapes[site] = site_shape
68         else:
69             if (return_sites is not None) or (site not in reshaped_samples):
70                 return_site_shapes[site] = site_shape
71 
72     # handle _RETURN site
73     if isinstance(return_sites, (list, tuple, set)) and '_RETURN' in return_sites:
74         value = model_trace.nodes['_RETURN']['value']
75         shape = (num_samples,) + value.shape if torch.is_tensor(value) else None
76         return_site_shapes['_RETURN'] = shape
77 
78     if not parallel:
79         return _predictive_sequential(model, posterior_samples, model_args, model_kwargs, num_samples,
80                                       return_site_shapes.keys(), return_trace=False)
81 
82     trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\
83         .get_trace(*model_args, **model_kwargs)
84     predictions = {}
85     for site, shape in return_site_shapes.items():
86         value = trace.nodes[site]['value']
87         if site == '_RETURN' and shape is None:
88             predictions[site] = value
89             continue
90         if value.numel() < reduce((lambda x, y: x * y), shape):
91             predictions[site] = value.expand(shape)
92         else:
93             predictions[site] = value.reshape(shape)
94 
95     return predictions
96 
97 
98 class Predictive:
99     """
100     This class is used to construct predictive distribution. The predictive distribution is obtained
101     by running model conditioned on latent samples from `posterior_samples`.
102 
103     .. warning::
104         The interface for the :class:`Predictive` class is experimental, and
105         might change in the future.
106 
107     :param model: Python callable containing Pyro primitives.
108     :param dict posterior_samples: dictionary of samples from the posterior.
109     :param callable guide: optional guide to get posterior samples of sites not present
110         in `posterior_samples`.
111     :param int num_samples: number of samples to draw from the predictive distribution.
112         This argument has no effect if ``posterior_samples`` is non-empty, in which case,
113         the leading dimension size of samples in ``posterior_samples`` is used.
114     :param return_sites: sites to return; by default only sample sites not present
115         in `posterior_samples` are returned.
116     :type return_sites: list, tuple, or set
117     :param bool parallel: predict in parallel by wrapping the existing model
118         in an outermost `plate` messenger. Note that this requires that the model has
119         all batch dims correctly annotated via :class:`~pyro.plate`. Default is `False`.
120     """
121     def __init__(self, model, posterior_samples=None, guide=None, num_samples=None,
122                  return_sites=None, parallel=False):
123         if posterior_samples is None and num_samples is None:
124             raise ValueError("Either posterior_samples or num_samples must be specified.")
125 
126         posterior_samples = {} if posterior_samples is None else posterior_samples
127 
128         for name, sample in posterior_samples.items():
129             batch_size = sample.shape[0]
130             if num_samples is None:
131                 num_samples = batch_size
132             elif num_samples != batch_size:
133                 warnings.warn("Sample's leading dimension size {} is different from the "
134                               "provided {} num_samples argument. Defaulting to {}."
135                               .format(batch_size, num_samples, batch_size), UserWarning)
136                 num_samples = batch_size
137 
138         if num_samples is None:
139             raise ValueError("No sample sites in posterior samples to infer `num_samples`.")
140 
141         if return_sites is not None:
142             assert isinstance(return_sites, (list, tuple, set))
143 
144         self.model = model
145         self.posterior_samples = {} if posterior_samples is None else posterior_samples
146         self.num_samples = num_samples
147         self.guide = guide
148         self.return_sites = return_sites
149         self.parallel = parallel
150 
151     def get_samples(self, *args, **kwargs):
152         """
153         Returns dict of samples from the predictive distribution. By default, only sample sites not
154         contained in `posterior_samples` are returned. This can be modified by changing the
155         `return_sites` keyword argument of this :class:`Predictive` instance.
156 
157         :param args: model arguments.
158         :param kwargs: model keyword arguments.
159         """
160         posterior_samples = self.posterior_samples
161         if self.guide is not None:
162             # use return_sites='' as a special signal to return all sites
163             posterior_samples = _predictive(self.guide, posterior_samples, self.num_samples,
164                                             return_sites='', parallel=self.parallel,
165                                             model_args=args, model_kwargs=kwargs)
166         return _predictive(self.model, posterior_samples, self.num_samples,
167                            return_sites=self.return_sites, parallel=self.parallel,
168                            model_args=args, model_kwargs=kwargs)
169 
170     def get_vectorized_trace(self, *args, **kwargs):
171         """
172         Returns a single vectorized `trace` from the predictive distribution. Note that this
173         requires that the model has all batch dims correctly annotated via :class:`~pyro.plate`.
174 
175         :param args: model arguments.
176         :param kwargs: model keyword arguments.
177         """
178         posterior_samples = self.posterior_samples
179         if self.guide is not None:
180             # use return_sites='' as a special signal to return all sites
181             posterior_samples = _predictive(self.guide, posterior_samples, self.num_samples,
182                                             return_sites='', parallel=self.parallel,
183                                             model_args=args, model_kwargs=kwargs)
184         return _predictive(self.model, posterior_samples, self.num_samples,
185                            return_trace=True, model_args=args, model_kwargs=kwargs)
186
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 _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 @@
         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 @@
 
     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 @@
 
     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)

{"golden_diff": "diff --git a/pyro/infer/predictive.py b/pyro/infer/predictive.py\n--- a/pyro/infer/predictive.py\n+++ b/pyro/infer/predictive.py\n@@ -28,7 +28,7 @@\n \n \n def _predictive_sequential(model, posterior_samples, model_args, model_kwargs,\n-                           num_samples, sample_sites, return_trace=False):\n+                           num_samples, return_site_shapes, return_trace=False):\n     collected = []\n     samples = [{k: v[i] for k, v in posterior_samples.items()} for i in range(num_samples)]\n     for i in range(num_samples):\n@@ -36,10 +36,13 @@\n         if return_trace:\n             collected.append(trace)\n         else:\n-            collected.append({site: trace.nodes[site]['value'] for site in sample_sites})\n+            collected.append({site: trace.nodes[site]['value'] for site in return_site_shapes})\n \n-    return collected if return_trace else {site: torch.stack([s[site] for s in collected])\n-                                           for site in sample_sites}\n+    if return_trace:\n+        return collected\n+    else:\n+        return {site: torch.stack([s[site] for s in collected]).reshape(shape)\n+                for site, shape in return_site_shapes.items()}\n \n \n def _predictive(model, posterior_samples, num_samples, return_sites=None,\n@@ -61,7 +64,8 @@\n \n     return_site_shapes = {}\n     for site in model_trace.stochastic_nodes + model_trace.observation_nodes:\n-        site_shape = (num_samples,) + model_trace.nodes[site]['value'].shape\n+        append_ndim = max_plate_nesting - len(model_trace.nodes[site][\"fn\"].batch_shape)\n+        site_shape = (num_samples,) + (1,) * append_ndim + model_trace.nodes[site]['value'].shape\n         if isinstance(return_sites, (list, tuple, set)):\n             if site in return_sites:\n                 return_site_shapes[site] = site_shape\n@@ -77,7 +81,7 @@\n \n     if not parallel:\n         return _predictive_sequential(model, posterior_samples, model_args, model_kwargs, num_samples,\n-                                      return_site_shapes.keys(), return_trace=False)\n+                                      return_site_shapes, return_trace=False)\n \n     trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\\\n         .get_trace(*model_args, **model_kwargs)\n", "issue": "Predictive class fails for non-scalar sites\nI usually use raw poutines to for prediction:\r\n```py\r\nvectorize = pyro.plate(\"_vectorize\", num_samples, dim=-2)\r\ntrace = poutine.trace(vectorize(guide)).get_trace()\r\nsamples = poutine.replay(vectorize(model), trace)()\r\n```\r\nWhen instead I try the new `Predictive` class, this fails on models with nontrivial plate structure:\r\n```py\r\npredictive = Predictive(model, guide=guide, return_sites=[\"x\", \"y\"],\r\n                        num_samples=num_samples, parallel=True)\r\nsamples = predictive.get_samples()\r\n```\r\nI 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.\r\n\r\nI have added a failing test in this branch. Feel free to take over the branch and create a PR:\r\nhttps://github.com/pyro-ppl/pyro/pull/new/fix-preditive-shapes\n", "before_files": [{"content": "from functools import reduce\nimport warnings\n\nimport torch\n\nimport pyro\nimport pyro.poutine as poutine\nfrom pyro.poutine.util import prune_subsample_sites\n\n\ndef _guess_max_plate_nesting(model, args, kwargs):\n    \"\"\"\n    Guesses max_plate_nesting by running the model once\n    without enumeration. This optimistically assumes static model\n    structure.\n    \"\"\"\n    with poutine.block():\n        model_trace = poutine.trace(model).get_trace(*args, **kwargs)\n    sites = [site for site in model_trace.nodes.values()\n             if site[\"type\"] == \"sample\"]\n\n    dims = [frame.dim\n            for site in sites\n            for frame in site[\"cond_indep_stack\"]\n            if frame.vectorized]\n    max_plate_nesting = -min(dims) if dims else 0\n    return max_plate_nesting\n\n\ndef _predictive_sequential(model, posterior_samples, model_args, model_kwargs,\n                           num_samples, sample_sites, return_trace=False):\n    collected = []\n    samples = [{k: v[i] for k, v in posterior_samples.items()} for i in range(num_samples)]\n    for i in range(num_samples):\n        trace = poutine.trace(poutine.condition(model, samples[i])).get_trace(*model_args, **model_kwargs)\n        if return_trace:\n            collected.append(trace)\n        else:\n            collected.append({site: trace.nodes[site]['value'] for site in sample_sites})\n\n    return collected if return_trace else {site: torch.stack([s[site] for s in collected])\n                                           for site in sample_sites}\n\n\ndef _predictive(model, posterior_samples, num_samples, return_sites=None,\n                return_trace=False, parallel=False, model_args=(), model_kwargs={}):\n    max_plate_nesting = _guess_max_plate_nesting(model, model_args, model_kwargs)\n    vectorize = pyro.plate(\"_num_predictive_samples\", num_samples, dim=-max_plate_nesting-1)\n    model_trace = prune_subsample_sites(poutine.trace(model).get_trace(*model_args, **model_kwargs))\n    reshaped_samples = {}\n\n    for name, sample in posterior_samples.items():\n        sample_shape = sample.shape[1:]\n        sample = sample.reshape((num_samples,) + (1,) * (max_plate_nesting - len(sample_shape)) + sample_shape)\n        reshaped_samples[name] = sample\n\n    if return_trace:\n        trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\\\n            .get_trace(*model_args, **model_kwargs)\n        return trace\n\n    return_site_shapes = {}\n    for site in model_trace.stochastic_nodes + model_trace.observation_nodes:\n        site_shape = (num_samples,) + model_trace.nodes[site]['value'].shape\n        if isinstance(return_sites, (list, tuple, set)):\n            if site in return_sites:\n                return_site_shapes[site] = site_shape\n        else:\n            if (return_sites is not None) or (site not in reshaped_samples):\n                return_site_shapes[site] = site_shape\n\n    # handle _RETURN site\n    if isinstance(return_sites, (list, tuple, set)) and '_RETURN' in return_sites:\n        value = model_trace.nodes['_RETURN']['value']\n        shape = (num_samples,) + value.shape if torch.is_tensor(value) else None\n        return_site_shapes['_RETURN'] = shape\n\n    if not parallel:\n        return _predictive_sequential(model, posterior_samples, model_args, model_kwargs, num_samples,\n                                      return_site_shapes.keys(), return_trace=False)\n\n    trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\\\n        .get_trace(*model_args, **model_kwargs)\n    predictions = {}\n    for site, shape in return_site_shapes.items():\n        value = trace.nodes[site]['value']\n        if site == '_RETURN' and shape is None:\n            predictions[site] = value\n            continue\n        if value.numel() < reduce((lambda x, y: x * y), shape):\n            predictions[site] = value.expand(shape)\n        else:\n            predictions[site] = value.reshape(shape)\n\n    return predictions\n\n\nclass Predictive:\n    \"\"\"\n    This class is used to construct predictive distribution. The predictive distribution is obtained\n    by running model conditioned on latent samples from `posterior_samples`.\n\n    .. warning::\n        The interface for the :class:`Predictive` class is experimental, and\n        might change in the future.\n\n    :param model: Python callable containing Pyro primitives.\n    :param dict posterior_samples: dictionary of samples from the posterior.\n    :param callable guide: optional guide to get posterior samples of sites not present\n        in `posterior_samples`.\n    :param int num_samples: number of samples to draw from the predictive distribution.\n        This argument has no effect if ``posterior_samples`` is non-empty, in which case,\n        the leading dimension size of samples in ``posterior_samples`` is used.\n    :param return_sites: sites to return; by default only sample sites not present\n        in `posterior_samples` are returned.\n    :type return_sites: list, tuple, or set\n    :param bool parallel: predict in parallel by wrapping the existing model\n        in an outermost `plate` messenger. Note that this requires that the model has\n        all batch dims correctly annotated via :class:`~pyro.plate`. Default is `False`.\n    \"\"\"\n    def __init__(self, model, posterior_samples=None, guide=None, num_samples=None,\n                 return_sites=None, parallel=False):\n        if posterior_samples is None and num_samples is None:\n            raise ValueError(\"Either posterior_samples or num_samples must be specified.\")\n\n        posterior_samples = {} if posterior_samples is None else posterior_samples\n\n        for name, sample in posterior_samples.items():\n            batch_size = sample.shape[0]\n            if num_samples is None:\n                num_samples = batch_size\n            elif num_samples != batch_size:\n                warnings.warn(\"Sample's leading dimension size {} is different from the \"\n                              \"provided {} num_samples argument. Defaulting to {}.\"\n                              .format(batch_size, num_samples, batch_size), UserWarning)\n                num_samples = batch_size\n\n        if num_samples is None:\n            raise ValueError(\"No sample sites in posterior samples to infer `num_samples`.\")\n\n        if return_sites is not None:\n            assert isinstance(return_sites, (list, tuple, set))\n\n        self.model = model\n        self.posterior_samples = {} if posterior_samples is None else posterior_samples\n        self.num_samples = num_samples\n        self.guide = guide\n        self.return_sites = return_sites\n        self.parallel = parallel\n\n    def get_samples(self, *args, **kwargs):\n        \"\"\"\n        Returns dict of samples from the predictive distribution. By default, only sample sites not\n        contained in `posterior_samples` are returned. This can be modified by changing the\n        `return_sites` keyword argument of this :class:`Predictive` instance.\n\n        :param args: model arguments.\n        :param kwargs: model keyword arguments.\n        \"\"\"\n        posterior_samples = self.posterior_samples\n        if self.guide is not None:\n            # use return_sites='' as a special signal to return all sites\n            posterior_samples = _predictive(self.guide, posterior_samples, self.num_samples,\n                                            return_sites='', parallel=self.parallel,\n                                            model_args=args, model_kwargs=kwargs)\n        return _predictive(self.model, posterior_samples, self.num_samples,\n                           return_sites=self.return_sites, parallel=self.parallel,\n                           model_args=args, model_kwargs=kwargs)\n\n    def get_vectorized_trace(self, *args, **kwargs):\n        \"\"\"\n        Returns a single vectorized `trace` from the predictive distribution. Note that this\n        requires that the model has all batch dims correctly annotated via :class:`~pyro.plate`.\n\n        :param args: model arguments.\n        :param kwargs: model keyword arguments.\n        \"\"\"\n        posterior_samples = self.posterior_samples\n        if self.guide is not None:\n            # use return_sites='' as a special signal to return all sites\n            posterior_samples = _predictive(self.guide, posterior_samples, self.num_samples,\n                                            return_sites='', parallel=self.parallel,\n                                            model_args=args, model_kwargs=kwargs)\n        return _predictive(self.model, posterior_samples, self.num_samples,\n                           return_trace=True, model_args=args, model_kwargs=kwargs)\n", "path": "pyro/infer/predictive.py"}], "after_files": [{"content": "from functools import reduce\nimport warnings\n\nimport torch\n\nimport pyro\nimport pyro.poutine as poutine\nfrom pyro.poutine.util import prune_subsample_sites\n\n\ndef _guess_max_plate_nesting(model, args, kwargs):\n    \"\"\"\n    Guesses max_plate_nesting by running the model once\n    without enumeration. This optimistically assumes static model\n    structure.\n    \"\"\"\n    with poutine.block():\n        model_trace = poutine.trace(model).get_trace(*args, **kwargs)\n    sites = [site for site in model_trace.nodes.values()\n             if site[\"type\"] == \"sample\"]\n\n    dims = [frame.dim\n            for site in sites\n            for frame in site[\"cond_indep_stack\"]\n            if frame.vectorized]\n    max_plate_nesting = -min(dims) if dims else 0\n    return max_plate_nesting\n\n\ndef _predictive_sequential(model, posterior_samples, model_args, model_kwargs,\n                           num_samples, return_site_shapes, return_trace=False):\n    collected = []\n    samples = [{k: v[i] for k, v in posterior_samples.items()} for i in range(num_samples)]\n    for i in range(num_samples):\n        trace = poutine.trace(poutine.condition(model, samples[i])).get_trace(*model_args, **model_kwargs)\n        if return_trace:\n            collected.append(trace)\n        else:\n            collected.append({site: trace.nodes[site]['value'] for site in return_site_shapes})\n\n    if return_trace:\n        return collected\n    else:\n        return {site: torch.stack([s[site] for s in collected]).reshape(shape)\n                for site, shape in return_site_shapes.items()}\n\n\ndef _predictive(model, posterior_samples, num_samples, return_sites=None,\n                return_trace=False, parallel=False, model_args=(), model_kwargs={}):\n    max_plate_nesting = _guess_max_plate_nesting(model, model_args, model_kwargs)\n    vectorize = pyro.plate(\"_num_predictive_samples\", num_samples, dim=-max_plate_nesting-1)\n    model_trace = prune_subsample_sites(poutine.trace(model).get_trace(*model_args, **model_kwargs))\n    reshaped_samples = {}\n\n    for name, sample in posterior_samples.items():\n        sample_shape = sample.shape[1:]\n        sample = sample.reshape((num_samples,) + (1,) * (max_plate_nesting - len(sample_shape)) + sample_shape)\n        reshaped_samples[name] = sample\n\n    if return_trace:\n        trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\\\n            .get_trace(*model_args, **model_kwargs)\n        return trace\n\n    return_site_shapes = {}\n    for site in model_trace.stochastic_nodes + model_trace.observation_nodes:\n        append_ndim = max_plate_nesting - len(model_trace.nodes[site][\"fn\"].batch_shape)\n        site_shape = (num_samples,) + (1,) * append_ndim + model_trace.nodes[site]['value'].shape\n        if isinstance(return_sites, (list, tuple, set)):\n            if site in return_sites:\n                return_site_shapes[site] = site_shape\n        else:\n            if (return_sites is not None) or (site not in reshaped_samples):\n                return_site_shapes[site] = site_shape\n\n    # handle _RETURN site\n    if isinstance(return_sites, (list, tuple, set)) and '_RETURN' in return_sites:\n        value = model_trace.nodes['_RETURN']['value']\n        shape = (num_samples,) + value.shape if torch.is_tensor(value) else None\n        return_site_shapes['_RETURN'] = shape\n\n    if not parallel:\n        return _predictive_sequential(model, posterior_samples, model_args, model_kwargs, num_samples,\n                                      return_site_shapes, return_trace=False)\n\n    trace = poutine.trace(poutine.condition(vectorize(model), reshaped_samples))\\\n        .get_trace(*model_args, **model_kwargs)\n    predictions = {}\n    for site, shape in return_site_shapes.items():\n        value = trace.nodes[site]['value']\n        if site == '_RETURN' and shape is None:\n            predictions[site] = value\n            continue\n        if value.numel() < reduce((lambda x, y: x * y), shape):\n            predictions[site] = value.expand(shape)\n        else:\n            predictions[site] = value.reshape(shape)\n\n    return predictions\n\n\nclass Predictive:\n    \"\"\"\n    This class is used to construct predictive distribution. The predictive distribution is obtained\n    by running model conditioned on latent samples from `posterior_samples`.\n\n    .. warning::\n        The interface for the :class:`Predictive` class is experimental, and\n        might change in the future.\n\n    :param model: Python callable containing Pyro primitives.\n    :param dict posterior_samples: dictionary of samples from the posterior.\n    :param callable guide: optional guide to get posterior samples of sites not present\n        in `posterior_samples`.\n    :param int num_samples: number of samples to draw from the predictive distribution.\n        This argument has no effect if ``posterior_samples`` is non-empty, in which case,\n        the leading dimension size of samples in ``posterior_samples`` is used.\n    :param return_sites: sites to return; by default only sample sites not present\n        in `posterior_samples` are returned.\n    :type return_sites: list, tuple, or set\n    :param bool parallel: predict in parallel by wrapping the existing model\n        in an outermost `plate` messenger. Note that this requires that the model has\n        all batch dims correctly annotated via :class:`~pyro.plate`. Default is `False`.\n    \"\"\"\n    def __init__(self, model, posterior_samples=None, guide=None, num_samples=None,\n                 return_sites=None, parallel=False):\n        if posterior_samples is None and num_samples is None:\n            raise ValueError(\"Either posterior_samples or num_samples must be specified.\")\n\n        posterior_samples = {} if posterior_samples is None else posterior_samples\n\n        for name, sample in posterior_samples.items():\n            batch_size = sample.shape[0]\n            if num_samples is None:\n                num_samples = batch_size\n            elif num_samples != batch_size:\n                warnings.warn(\"Sample's leading dimension size {} is different from the \"\n                              \"provided {} num_samples argument. Defaulting to {}.\"\n                              .format(batch_size, num_samples, batch_size), UserWarning)\n                num_samples = batch_size\n\n        if num_samples is None:\n            raise ValueError(\"No sample sites in posterior samples to infer `num_samples`.\")\n\n        if return_sites is not None:\n            assert isinstance(return_sites, (list, tuple, set))\n\n        self.model = model\n        self.posterior_samples = {} if posterior_samples is None else posterior_samples\n        self.num_samples = num_samples\n        self.guide = guide\n        self.return_sites = return_sites\n        self.parallel = parallel\n\n    def get_samples(self, *args, **kwargs):\n        \"\"\"\n        Returns dict of samples from the predictive distribution. By default, only sample sites not\n        contained in `posterior_samples` are returned. This can be modified by changing the\n        `return_sites` keyword argument of this :class:`Predictive` instance.\n\n        :param args: model arguments.\n        :param kwargs: model keyword arguments.\n        \"\"\"\n        posterior_samples = self.posterior_samples\n        if self.guide is not None:\n            # use return_sites='' as a special signal to return all sites\n            posterior_samples = _predictive(self.guide, posterior_samples, self.num_samples,\n                                            return_sites='', parallel=self.parallel,\n                                            model_args=args, model_kwargs=kwargs)\n        return _predictive(self.model, posterior_samples, self.num_samples,\n                           return_sites=self.return_sites, parallel=self.parallel,\n                           model_args=args, model_kwargs=kwargs)\n\n    def get_vectorized_trace(self, *args, **kwargs):\n        \"\"\"\n        Returns a single vectorized `trace` from the predictive distribution. Note that this\n        requires that the model has all batch dims correctly annotated via :class:`~pyro.plate`.\n\n        :param args: model arguments.\n        :param kwargs: model keyword arguments.\n        \"\"\"\n        posterior_samples = self.posterior_samples\n        if self.guide is not None:\n            # use return_sites='' as a special signal to return all sites\n            posterior_samples = _predictive(self.guide, posterior_samples, self.num_samples,\n                                            return_sites='', parallel=self.parallel,\n                                            model_args=args, model_kwargs=kwargs)\n        return _predictive(self.model, posterior_samples, self.num_samples,\n                           return_trace=True, model_args=args, model_kwargs=kwargs)\n", "path": "pyro/infer/predictive.py"}]}