code: def potential_grad(potential_fn, z):
      """
      Gradient of `potential_fn` w.r.t. parameters z.
  
      :param potential_fn: python callable that takes in a dictionary of parameters
          and returns the potential energy.
      :param dict z: dictionary of parameter values keyed by site name.
      :return: tuple of `(z_grads, potential_energy)`, where `z_grads` is a dictionary
          with the same keys as `z` containing gradients and potential_energy is a
          torch scalar.
      """
      z_keys, z_nodes = zip(*z.items())
      for node in z_nodes:
          node.requires_grad_(True)
      try:
          potential_energy = potential_fn(z)
      # handle exceptions as defined in the exception registry
      except Exception as e:
          if any(h(e) for h in _EXCEPTION_HANDLERS.values()):
              grads = {k: v.new_zeros(v.shape) for k, v in z.items()}
              return grads, z_nodes[0].new_tensor(float("nan"))
          else:
              raise e
      grads = grad(potential_energy, z_nodes)
      for node in z_nodes:
          node.requires_grad_(False)
      return dict(zip(z_keys, grads)), potential_energy.detach()
  
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 3 for calling a global function with example usage: potential_grad(potential_fn, z) and returns: dict,  potential_energy, grads,  z_nodes