code: def velocity_verlet(
      z,
      r,
      potential_fn,
      kinetic_grad,
      step_size,
      num_steps=1,
      z_grads=None,
  ):
      r"""
      Second order symplectic integrator that uses the velocity verlet algorithm.
  
      :param dict z: dictionary of sample site names and their current values
          (type :class:`~torch.Tensor`).
      :param dict r: dictionary of sample site names and corresponding momenta
          (type :class:`~torch.Tensor`).
      :param callable potential_fn: function that returns potential energy given z
          for each sample site. The negative gradient of the function with respect
          to ``z`` determines the rate of change of the corresponding sites'
          momenta ``r``.
      :param callable kinetic_grad: a function calculating gradient of kinetic energy
          w.r.t. momentum variable.
      :param float step_size: step size for each time step iteration.
      :param int num_steps: number of discrete time steps over which to integrate.
      :param torch.Tensor z_grads: optional gradients of potential energy at current ``z``.
      :return tuple (z_next, r_next, z_grads, potential_energy): next position and momenta,
          together with the potential energy and its gradient w.r.t. ``z_next``.
      """
      z_next = z.copy()
      r_next = r.copy()
      for _ in range(num_steps):
          (
              z_next,
              r_next,
              z_grads,
              potential_energy,
          ) = _single_step_verlet(
              z_next,
              r_next,
              potential_fn,
              kinetic_grad,
              step_size,
              z_grads,
          )
      return z_next, r_next, z_grads, potential_energy
  
file_path: pyro/pyro/ops/integrator.py
imports: import warnings
  from typing import Callable, Dict
  from torch.autograd import grad
label_desc: Use image node_id 1 for calling a global function with example usage: velocity_verlet(z, r, potential_fn, kinetic_grad, step_size, num_steps, z_grads) and returns: z_next,  r_next,  z_grads,  potential_energy