class_docstr: A basic implementation of Stein Variational Gradient Descent as described in reference [1].
  
  :param model: The model (callable containing Pyro primitives). Model must be fully vectorized
      and may only contain continuous latent variables.
  :param kernel: a SVGD compatible kernel like :class:`RBFSteinKernel`.
  :param optim: A wrapper for a PyTorch optimizer.
  :type optim: pyro.optim.PyroOptim
  :param int num_particles: The number of particles used in SVGD.
  :param int max_plate_nesting: The max number of nested :func:`pyro.plate` contexts in the model.
  :param str mode: Whether to use a Kernelized Stein Discrepancy that makes use of `multivariate`
      test functions (as in [1]) or `univariate` test functions (as in [2]). Defaults to `univariate`.
  
  Example usage:
  
  .. code-block:: python
  
      from pyro.infer import SVGD, RBFSteinKernel
      from pyro.optim import Adam
  
      kernel = RBFSteinKernel()
      adam = Adam({"lr": 0.1})
      svgd = SVGD(model, kernel, adam, num_particles=50, max_plate_nesting=0)
  
      for step in range(500):
          svgd.step(model_arg1, model_arg2)
  
      final_particles = svgd.get_named_particles()
  
  References
  
  [1] "Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm,"
      Qiang Liu, Dilin Wang
  [2] "Kernelized Complete Conditional Stein Discrepancy,"
      Raghav Singhal, Saad Lahlou, Rajesh Ranganath
docstr: Computes the SVGD gradient, passing args and kwargs to the model,
  and takes a gradient step.
  
  :return dict: A dictionary of the form {name: float}, where each float
      is a mean squared gradient. This can be used to monitor the convergence of SVGD.
docstr_tok: ["Computes","the","SVGD","gradient",",","passing","args","and","kwargs","to","the","model",",","and","takes","a","gradient","step",".",":","return","dict",":","A","dictionary","of","the","form","{","name",":","float","}",",","where","each","float","is","a","mean","squared","gradient",".","This","can","be","used","to","monitor","the","convergence","of","SVGD","."]
code: def step(self, *args, **kwargs):
      """
      Computes the SVGD gradient, passing args and kwargs to the model,
      and takes a gradient step.
  
      :return dict: A dictionary of the form {name: float}, where each float
          is a mean squared gradient. This can be used to monitor the convergence of SVGD.
      """
      # compute gradients of log model joint
      with torch.enable_grad(), poutine.trace(
          param_only=True
      ) as param_capture:
          loss = self.loss(self.model, self.guide, *args, **kwargs)
          loss.backward()
  
      # get particles used in the _SVGDGuide and reshape to have num_particles leading dimension
      particles = pyro.param("svgd_particles").unconstrained()
      reshaped_particles = particles.reshape(self.num_particles, -1)
      reshaped_particles_grad = particles.grad.reshape(
          self.num_particles, -1
      )
  
      # compute kernel ingredients
      log_kernel, kernel_grad = self.kernel.log_kernel_and_grad(
          reshaped_particles
      )
  
      if self.mode == "multivariate":
          kernel = log_kernel.sum(-1).exp()
          assert kernel.shape == (
              self.num_particles,
              self.num_particles,
          )
          attractive_grad = torch.mm(kernel, reshaped_particles_grad)
          repulsive_grad = torch.einsum(
              "nm,nm...->n...", kernel, kernel_grad
          )
      elif self.mode == "univariate":
          kernel = log_kernel.exp()
          assert kernel.shape == (
              self.num_particles,
              self.num_particles,
              reshaped_particles.size(-1),
          )
          attractive_grad = torch.einsum(
              "nmd,md->nd", kernel, reshaped_particles_grad
          )
          repulsive_grad = torch.einsum(
              "nmd,nmd->nd", kernel, kernel_grad
          )
  
      # combine the attractive and repulsive terms in the SVGD gradient
      assert attractive_grad.shape == repulsive_grad.shape
      particles.grad = (attractive_grad + repulsive_grad).reshape(
          particles.shape
      ) / self.num_particles
  
      # compute per-parameter mean squared gradients
      squared_gradients = {
          site["name"]: value.mean().item()
          for site, value in self.guide._unpack_latent(
              particles.grad.pow(2.0)
          )
      }
  
      # torch.optim objects gets instantiated for any params that haven't been seen yet
      params = set(
          site["value"].unconstrained()
          for site in param_capture.trace.nodes.values()
      )
      self.optim(params)
  
      # zero gradients
      pyro.infer.util.zero_grads(params)
  
      # return per-parameter mean squared gradients to user
      return squared_gradients
  
imports: import math
  from abc import ABCMeta, abstractmethod
  import torch
  from torch.distributions import biject_to
  import pyro
  from pyro import poutine
  from pyro.distributions import Delta
  from pyro.distributions.util import copy_docs_from
  from pyro.infer.autoguide.guides import AutoContinuous
  from pyro.infer.autoguide.initialization import init_to_sample
  from pyro.infer.trace_elbo import Trace_ELBO
label_desc: Use image node_id 3 for calling the SVGD obj's underlying member method code with example usage: obj.step() and returns: squared_gradients