code: def _jit_adjoint(equation, *operands, **kwargs):
      """
      Runs einsum in forward-backward mode using ``pyro.ops.adjoint``.
  
      This simulates serving predictions from an undirected graphical model.
      """
      backend = kwargs.pop("backend", "pyro.ops.einsum.torch_sample")
      key = (
          backend,
          equation,
          tuple(x.shape for x in operands),
          kwargs["plates"],
      )
      if key not in _CACHE:
          # This wraps a complete adjoint algorithm call.
          @ignore_jit_warnings()
          def _forward_backward(*operands):
              # First we request backward results on each input operand.
              # This is the pyro.ops.adjoint equivalent of torch's .requires_grad_().
              for operand in operands:
                  require_backward(operand)
  
              # Next we run the forward pass.
              results = einsum(
                  equation, *operands, backend=backend, **kwargs
              )
  
              # The we run a backward pass.
              for result in results:
                  result._pyro_backward()
  
              # Finally we retrieve results from the ._pyro_backward_result attribute
              # that has been set on each input operand. If you only want results on a
              # subset of operands, you can call require_backward() on only those.
              results = []
              for x in operands:
                  results.append(x._pyro_backward_result)
                  x._pyro_backward_result = None
  
              return tuple(results)
  
          _CACHE[key] = torch.jit.trace(
              _forward_backward, operands, check_trace=False
          )
  
      return _CACHE[key](*operands)
  
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 4 for calling a global function with example usage: _jit_adjoint(equation) and returns: _CACHE, tuple