code: def main(args):
      if args.cuda:
          torch.set_default_device("cuda")
  
      if args.method == "all":
          for method in [
              "prob",
              "logprob",
              "gradient",
              "marginal",
              "map",
              "sample",
          ]:
              args.method = method
              main(args)
          return
  
      print(
          "Plate size  Time per iteration of {} (ms)".format(
              args.method
          )
      )
      fn = globals()["jit_{}".format(args.method)]
      equation = args.equation
      plates = args.plates
      inputs, outputs = equation.split("->")
      inputs = inputs.split(",")
  
      # Vary all plate sizes at the same time.
      for plate_size in range(8, 1 + args.max_plate_size, 8):
          operands = []
          for dims in inputs:
              shape = torch.Size(
                  [
                      plate_size if d in plates else args.dim_size
                      for d in dims
                  ]
              )
              operands.append(
                  (torch.empty(shape).uniform_() + 0.5).requires_grad_()
              )
  
          time = time_fn(
              fn,
              equation,
              *operands,
              plates=plates,
              modulo_total=True,
              iters=args.iters,
          )
          print(
              "{: <11s} {:0.4g}".format(
                  "{} ** {}".format(plate_size, len(args.plates)),
                  time * 1e3,
              )
          )
  
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 9 for calling a global function with example usage: main(args) without return types