code: def testPHMC(self):
      step_size = self._constant(0.2)
      num_steps = 2000
      num_leapfrog_steps = 10
      state = jnp.ones([16, 2], dtype=self._dtype)
  
      base_mean = self._constant([2.0, 3.0])
      base_scale = self._constant([2.0, 0.5])
  
      def target_log_prob_fn(x):
          return (
              -jnp.sum(
                  0.5 * jnp.square((x - base_mean) / base_scale), -1
              ),
              (),
          )
  
      def kernel(phmc_state, seed):
          phmc_seed, seed = util.split_seed(seed, 2)
          (
              phmc_state,
              _,
          ) = prefab.persistent_hamiltonian_monte_carlo_step(
              phmc_state,
              step_size=step_size,
              num_integrator_steps=num_leapfrog_steps,
              target_log_prob_fn=target_log_prob_fn,
              noise_fraction=self._constant(0.5),
              mh_drift=self._constant(0.127),
              seed=phmc_seed,
          )
          return (phmc_state, seed), phmc_state.state
  
      seed = self._make_seed(_test_seed())
  
      _, chain = jax.jit(
          lambda state, seed: fun_mc.trace(  # pylint: disable=g-long-lambda
              state=(
                  prefab.persistent_hamiltonian_monte_carlo_init(
                      state, target_log_prob_fn
                  ),
                  seed,
              ),
              fn=kernel,
              num_steps=num_steps,
          )
      )(state, seed)
      # Discard the warmup samples.
      chain = chain[1000:]
  
      sample_mean = jnp.mean(chain, axis=[0, 1])
      sample_var = jnp.var(chain, axis=[0, 1])
  
      true_samples = (
          util.random_normal(
              shape=[4096, 2], dtype=self._dtype, seed=seed
          )
          * base_scale
          + base_mean
      )
  
      true_mean = jnp.mean(true_samples, axis=0)
      true_var = jnp.var(true_samples, axis=0)
  
      self.assertAllClose(true_mean, sample_mean, rtol=0.1, atol=0.1)
      self.assertAllClose(true_var, sample_var, rtol=0.1, atol=0.1)
  
imports: import functools
  import os
  import jax
  from jax import config
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import prefab
  from fun_mc import test_util
  from fun_mc import util_tfp
label_desc: Use image node_id 9 for calling the PrefabTest obj's underlying member method code with example usage: obj.testPHMC() and returns: unknown,  , phmc_state,  seed,  phmc_state