code: 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()
      assert set(actual) == set(expected)
      assert actual["x"].shape == expected["x"].shape
      assert actual["y"].shape == expected["y"].shape
  
file_path: pyro/tests/infer/test_predictive.py
imports: import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.optim
  import pyro.poutine
  from pyro.infer import SVI, Predictive, Trace_ELBO
  from pyro.infer.autoguide import AutoDelta, AutoDiagonalNormal
  from tests.common import assert_close
label_desc: Use image node_id 8 for calling a global function with example usage: test_shapes(parallel) and returns: dict