code: def benchmark_eight_schools_hmc(
      num_results=int(5e3),
      num_burnin_steps=int(3e3),
      num_leapfrog_steps=3,
      step_size=0.4,
  ):
      """Runs HMC on the eight-schools unnormalized posterior."""
  
      num_schools = 8
      treatment_effects = tf.constant(
          [28, 8, -3, 7, -1, 1, 18, 12],
          dtype=np.float32,
          name="treatment_effects",
      )
      treatment_stddevs = tf.constant(
          [15, 10, 16, 11, 9, 11, 10, 18],
          dtype=np.float32,
          name="treatment_stddevs",
      )
  
      def unnormalized_posterior_log_prob(
          avg_effect, avg_stddev, school_effects_standard
      ):
          """Eight-schools unnormalized log posterior."""
          return eight_schools_joint_log_prob(
              treatment_effects,
              treatment_stddevs,
              avg_effect,
              avg_stddev,
              school_effects_standard,
          )
  
      if tf.executing_eagerly():
          sample_chain = tf.function(sample.sample_chain)
      else:
          sample_chain = sample.sample_chain
  
      def computation():
          """The benchmark computation."""
          _, kernel_results = sample_chain(
              num_results=num_results,
              num_burnin_steps=num_burnin_steps,
              current_state=(
                  tf.zeros([], name="init_avg_effect"),
                  tf.zeros([], name="init_avg_stddev"),
                  tf.ones(
                      [num_schools], name="init_school_effects_standard"
                  ),
              ),
              kernel=hmc.HamiltonianMonteCarlo(
                  target_log_prob_fn=unnormalized_posterior_log_prob,
                  step_size=step_size,
                  num_leapfrog_steps=num_leapfrog_steps,
              ),
          )
  
          return kernel_results.is_accepted
  
      # Let's force evaluation of graph to ensure build time is not part of our time
      # trial.
      is_accepted_tensor = computation()
      if not tf.executing_eagerly():
          session = tf1.Session()
          session.run(is_accepted_tensor)
  
      start_time = time.time()
      if tf.executing_eagerly():
          is_accepted = computation()
      else:
          is_accepted = session.run(is_accepted_tensor)
      wall_time = time.time() - start_time
  
      num_accepted = np.sum(is_accepted)
      acceptance_rate = np.float32(num_accepted) / np.float32(
          num_results
      )
  
      return dict(
          iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
          extras={"acceptance_rate": acceptance_rate},
          wall_time=wall_time,
      )
  
file_path: probability/tensorflow_probability/python/mcmc/eight_schools_hmc.py
imports: import time
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 3 for calling a global function with example usage: benchmark_eight_schools_hmc(num_results, num_burnin_steps, num_leapfrog_steps, step_size) and returns: dict, eight_schools_joint_log_prob, kernel_results