code: def test_shapes(shape, stein_kernel):
      pyro.clear_param_store()
      shape1, shape2 = (5,) + shape, shape + (6,)
  
      mean_init1 = (
          torch.arange(_product(shape1)).double().reshape(shape1)
          / 100.0
      )
      mean_init2 = (
          torch.arange(_product(shape2)).double().reshape(shape2)
      )
  
      def model():
          pyro.sample(
              "z1",
              dist.LogNormal(mean_init1, 1.0e-8).to_event(len(shape1)),
          )
          pyro.sample("scalar", dist.Normal(0.0, 1.0))
          pyro.sample(
              "z2",
              dist.Normal(mean_init2, 1.0e-8).to_event(len(shape2)),
          )
  
      num_particles = 7
      svgd = SVGD(
          model, stein_kernel(), Adam({"lr": 0.0}), num_particles, 0
      )
  
      for step in range(2):
          svgd.step()
  
      particles = svgd.get_named_particles()
      assert particles["z1"].shape == (num_particles,) + shape1
      assert particles["z2"].shape == (num_particles,) + shape2
  
      for particle in range(num_particles):
          assert_equal(
              particles["z1"][particle, ...],
              mean_init1.exp(),
              prec=1.0e-6,
          )
          assert_equal(
              particles["z2"][particle, ...], mean_init2, prec=1.0e-6
          )
  
file_path: pyro/tests/infer/test_svgd.py
imports: import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.infer import SVGD, IMQSteinKernel, RBFSteinKernel
  from pyro.infer.autoguide.utils import _product
  from pyro.optim import Adam
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_shapes(shape, stein_kernel) without return types