code: def jit_prob(equation, *operands, **kwargs):
      """
      Runs einsum to compute the partition function.
  
      This is cheap but less numerically stable than using the torch_log backend.
      """
      key = "prob", equation, kwargs["plates"]
      if key not in _CACHE:
          # This simply wraps einsum for jit compilation.
          def _einsum(*operands):
              return einsum(equation, *operands, **kwargs)
  
          _CACHE[key] = torch.jit.trace(
              _einsum, 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 1 for calling a global function with example usage: jit_prob(equation) and returns: _CACHE, einsum