class_docstr: More careful tests that sampling/preconditioning is actually working.
class_docstr_tok: ["More","careful","tests","that","sampling\/preconditioning","is","actually","working","."]
code: def test_sets_kinetic_energy(self):
      dist = mvn_diag.MultivariateNormalDiag(
          scale_diag=tf.constant([0.1, 10.0])
      )
      step_size = 0.1
      kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=dist.log_prob,
          step_size=step_size,
          num_leapfrog_steps=1,
          store_parameters_in_results=True,
      )
      init_state = tf.constant([0.1, 0.1])
      kr = kernel.bootstrap_results(init_state)
  
      # Manually set the momentum distribution.
      kr = unnest.replace_innermost(kr, momentum_distribution=dist)
  
      # Take one leapfrog step using the kernel.
      _, nkr = kernel.one_step(
          init_state, kr, seed=test_util.test_seed()
      )
      # Need to evaluate here for consistency in graph mode.
      (
          momentum_parts,
          target_grad_parts,
          proposed_state,
          final_momentum,
          target_log_prob,
          grads_target_log_prob,
      ) = self.evaluate(
          [
              nkr.proposed_results.initial_momentum,
              nkr.accepted_results.grads_target_log_prob,
              nkr.proposed_state,
              nkr.proposed_results.final_momentum,
              nkr.proposed_results.target_log_prob,
              nkr.proposed_results.grads_target_log_prob,
          ]
      )
  
      # Take one leapfrog step manually.
      leapfrog = leapfrog_integrator.SimpleLeapfrogIntegrator(
          target_fn=dist.log_prob, step_sizes=[step_size], num_steps=1
      )
      # Again, need to evaluate here for graph mode consistency.
      (
          next_momentum,
          next_state,
          next_target_log_prob,
          grads_next_target_log_prob,
      ) = self.evaluate(
          leapfrog(
              momentum_parts=momentum_parts,
              state_parts=[init_state],
              target=dist.log_prob(init_state),
              target_grad_parts=target_grad_parts,
              kinetic_energy_fn=lambda x: -dist.log_prob(x),
          )
      )
  
      # Verify resulting states are the same
      self.assertAllClose(proposed_state, next_state[0])
      self.assertAllClose(final_momentum, next_momentum)
      self.assertAllClose(target_log_prob, next_target_log_prob)
      self.assertAllClose(
          grads_target_log_prob, grads_next_target_log_prob
      )
  
imports: import collections
  from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.distributions import distribution
  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 mvn_linear_operator
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.internal import auto_composite_tensor
  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.internal import unnest
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 6 for calling the PreconditionedHMCCorrectnessTest obj's underlying member method code with example usage: obj.test_sets_kinetic_energy() without return types