code: def jit_gradient(equation, *operands, **kwargs):
      """
      Runs einsum and calls backward on the partition function.
  
      This is simulates training an undirected graphical model.
      """
      key = "gradient", equation, kwargs["plates"]
      if key not in _CACHE:
          # This wraps einsum for jit compilation, but we will call backward on the result.
          def _einsum(*operands):
              return einsum(
                  equation,
                  *operands,
                  backend="pyro.ops.einsum.torch_log",
                  **kwargs,
              )
  
          _CACHE[key] = torch.jit.trace(
              _einsum, operands, check_trace=False
          )
  
      # Run forward pass.
      losses = _CACHE[key](*operands)
  
      # Work around PyTorch 1.0.0 bug https://github.com/pytorch/pytorch/issues/14875
      # whereby tuples of length 1 are unwrapped by the jit.
      if not isinstance(losses, tuple):
          losses = (losses,)
  
      # Run backward pass.
      grads = tuple(
          grad(loss, operands, retain_graph=True, allow_unused=True)
          for loss in losses
      )
      return grads
  
file_path: pyro/examples/einsum.py
imports: import argparse
  import timeit
  import torch
  from torch.autograd import grad
  from pyro.ops.contract import einsum
  from pyro.ops.einsum.adjoint import require_backward
  from pyro.util import ignore_jit_warnings
label_desc: Use image node_id 3 for calling a global function with example usage: jit_gradient(equation) and returns: grads, einsum