code: def _infer_hmc(args, data, model, init_values={}):
      logging.info("Running inference...")
      kernel = NUTS(
          model,
          full_mass=[("R0", "rho")],
          max_tree_depth=args.max_tree_depth,
          init_strategy=init_to_value(values=init_values),
          jit_compile=args.jit,
          ignore_jit_warnings=True,
      )
  
      # We'll define a hook_fn to log potential energy values during inference.
      # This is helpful to diagnose whether the chain is mixing.
      energies = []
  
      def hook_fn(kernel, *unused):
          e = float(kernel._potential_energy_last)
          energies.append(e)
          if args.verbose:
              logging.info("potential = {:0.6g}".format(e))
  
      mcmc = MCMC(
          kernel,
          hook_fn=hook_fn,
          num_samples=args.num_samples,
          warmup_steps=args.warmup_steps,
      )
      mcmc.run(args, data)
      mcmc.summary()
      if args.plot:
          import matplotlib.pyplot as plt
  
          plt.figure(figsize=(6, 3))
          plt.plot(energies)
          plt.xlabel("MCMC step")
          plt.ylabel("potential energy")
          plt.title("MCMC energy trace")
          plt.tight_layout()
  
      samples = mcmc.get_samples()
      return samples
  
imports: import argparse
  import logging
  import math
  import re
  from collections import OrderedDict
  import torch
  import pyro
  import pyro.distributions
  import pyro.distributions.hmm
  import pyro.poutine
  from pyro.infer import MCMC, NUTS, config_enumerate, infer_discrete
  from pyro.infer.autoguide import init_to_value
  from pyro.ops.special import safe_log
  from pyro.ops.tensor_utils import convolve
  from pyro.util import warn_if_nan
label_desc: Use image node_id 6 for calling a global function with example usage: _infer_hmc(args, data, model, init_values) and returns: samples