code: def run_hmc_on_model(
      model,
      num_chains,
      num_steps,
      num_leapfrog_steps,
      step_size,
      target_accept_prob=0.9,
      seed=None,
      dtype=tf.float32,
      use_xla=True,
  ):
      """Runs HMC on a target.
  
      Args:
        model: The model to validate.
        num_chains: Number of chains to run in parallel.
        num_steps: Total number of steps to take. The first half are used to warm up
          the sampler.
        num_leapfrog_steps: Number of leapfrog steps to take.
        step_size: Step size to use.
        target_accept_prob: Target acceptance probability.
        seed: Optional seed to use. By default, `test_util.test_seed()` is used.
        dtype: DType to use for the algorithm.
        use_xla: Whether to use XLA.
  
      Returns:
        mcmc_results: `MCMCResults`.
      """
      step_size = tf.convert_to_tensor(step_size, dtype)
  
      def target_log_prob_fn(*x):
          x = tf.nest.pack_sequence_as(model.dtype, x)
          return model.unnormalized_log_prob(x)
  
      if seed is None:
          seed = test_util.test_seed(sampler_type="stateless")
      current_state = tf.nest.map_structure(
          lambda b, e: b(  # pylint: disable=g-long-lambda
              tf.zeros([num_chains] + list(e), dtype=dtype)
          ),
          model.default_event_space_bijector,
          model.event_shape,
      )
  
      # tfp.mcmc only works well with lists.
      current_state = tf.nest.flatten(current_state)
  
      hmc = tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          num_leapfrog_steps=num_leapfrog_steps,
          step_size=[
              tf.fill(s.shape, step_size) for s in current_state
          ],
      )
      hmc = tfp.mcmc.TransformedTransitionKernel(
          hmc, tf.nest.flatten(model.default_event_space_bijector)
      )
      hmc = tfp.mcmc.DualAveragingStepSizeAdaptation(
          hmc,
          num_adaptation_steps=int(num_steps // 2 * 0.8),
          target_accept_prob=target_accept_prob,
      )
  
      # Subtle: Under JAX, there needs to be a data dependency on the input for
      # jitting to work.
      chain, is_accepted = tf.function(
          lambda current_state: tfp.mcmc.sample_chain(  # pylint: disable=g-long-lambda
              current_state=current_state,
              kernel=hmc,
              num_results=num_steps // 2,
              num_burnin_steps=num_steps // 2,
              trace_fn=lambda _, pkr: (  # pylint: disable=g-long-lambda
                  pkr.inner_results.inner_results.is_accepted
              ),
              seed=seed,
          ),
          autograph=False,
          jit_compile=use_xla,
      )(current_state)
  
      accept_rate = tf.reduce_mean(tf.cast(is_accepted, dtype))
      ess = tf.nest.map_structure(
          lambda c: tfp.mcmc.effective_sample_size(  # pylint: disable=g-long-lambda
              c, cross_chain_dims=1, filter_beyond_positive_pairs=True
          ),
          chain,
      )
      r_hat = tf.nest.map_structure(
          tfp.mcmc.potential_scale_reduction, chain
      )
  
      mcmc_results = MCMCResults(
          chain=tf.nest.pack_sequence_as(
              model.default_event_space_bijector, chain
          ),
          accept_rate=accept_rate,
          ess=ess,
          r_hat=r_hat,
      )
      return mcmc_results
  
imports: import collections
  import importlib
  from absl import flags
  from absl import logging
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 6 for calling a global function with example usage: run_hmc_on_model(model, num_chains, num_steps, num_leapfrog_steps, step_size, target_accept_prob, seed, dtype, use_xla) and returns: mcmc_results, model