code: def _run_hmc(self, num_results, initial_running_variance):
      mvn = mvn_diag.MultivariateNormalDiag(
          self.mvn_mean, scale_diag=self.mvn_scale
      )
      target_log_prob_fn = mvn.log_prob
      kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          num_leapfrog_steps=32,
          step_size=0.001,
      )
      kernel = dmma.DiagonalMassMatrixAdaptation(
          inner_kernel=kernel,
          initial_running_variance=initial_running_variance,
      )
  
      @tf.function
      def do_sample(seed):
          def trace_fn(_, pkr):
              return (
                  pkr.running_variance,
                  pkr.inner_results.accepted_results.momentum_distribution,
              )
  
          sample_seed, chain_seed = samplers.split_seed(seed, 2)
          draws, (rv_state, dist) = sample.sample_chain(
              num_results=num_results,
              current_state=tf.zeros(3),
              kernel=kernel,
              seed=chain_seed,
              trace_fn=trace_fn,
          )
  
          # sample_distributions returns `[dists], [samples]`, so the 0th
          # distribution corresponds to the 0th, and only, state part
          # The distribution is a BatchBroadcast containing a transformed
          # distribution, so we need to use .distribution.distribution
          batched_transformed_dist = dist.sample_distributions(
              seed=sample_seed
          )[0][0]
          momentum_dist = (
              batched_transformed_dist.distribution.distribution
          )
          final_precision_factor = momentum_dist.precision_factor.diag[
              -1
          ]
          # Evaluate here so we can check the value twice later
          final_precision = momentum_dist.precision.diag[-1]
          final_mean = rv_state[0].mean[-1]
          empirical_mean = tf.reduce_mean(draws, axis=0)
          # The final_precision is taken directly from the momentum distribution,
          # which never "sees" the last sample.
          empirical_variance = tf.math.reduce_variance(
              draws[:-1], axis=0
          )
          return RunHMCResults(
              draws=draws,
              final_mean=final_mean,
              final_precision_factor=final_precision_factor,
              final_precision=final_precision,
              empirical_mean=empirical_mean,
              empirical_variance=empirical_variance,
              true_mean=mvn.mean(),
              true_variance=mvn.variance(),
          )
  
      return self.evaluate(do_sample(test_util.test_seed()))
  
imports: import collections
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 3 for calling the DiagonalMassMatrixAdaptationTest obj's underlying member method code with example usage: obj._run_hmc(num_results, initial_running_variance) and returns: self, RunHMCResults, pkr,  pkr