class_docstr: A TestCase mixin for common tests on inference gym targets.
class_docstr_tok: ["A","TestCase","mixin","for","common","tests","on","inference","gym","targets","."]
docstr: Validates the ground truth of a model using HMC.
  
  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.
    mean_fudge_atol: Scalar `Tensor`. Additional atol to use when doing the
      Z-test for the sample mean.
    standard_deviation_fudge_atol: Scalar `Tensor`. Additional atol to use
      when doing the Z-test for the sample standard deviation,
    use_xla: Whether to use XLA.
docstr_tok: ["Validates","the","ground","truth","of","a","model","using","HMC",".","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",".","mean_fudge_atol",":","Scalar","`","Tensor","`",".","Additional","atol","to","use","when","doing","the","Z-test","for","the","sample","mean",".","standard_deviation_fudge_atol",":","Scalar","`","Tensor","`",".","Additional","atol","to","use","when","doing","the","Z-test","for","the","sample","standard","deviation",",","use_xla",":","Whether","to","use","XLA","."]
code: def validate_ground_truth_using_hmc(
      self,
      model,
      num_chains,
      num_steps,
      num_leapfrog_steps,
      step_size,
      target_accept_prob=0.9,
      seed=None,
      dtype=tf.float32,
      mean_fudge_atol=0,
      standard_deviation_fudge_atol=0,
      use_xla=True,
  ):
      """Validates the ground truth of a model using HMC.
  
      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.
        mean_fudge_atol: Scalar `Tensor`. Additional atol to use when doing the
          Z-test for the sample mean.
        standard_deviation_fudge_atol: Scalar `Tensor`. Additional atol to use
          when doing the Z-test for the sample standard deviation,
        use_xla: Whether to use XLA.
      """
      mcmc_results = self.evaluate(
          run_hmc_on_model(
              model,
              num_chains=num_chains,
              num_steps=num_steps,
              num_leapfrog_steps=num_leapfrog_steps,
              step_size=step_size,
              target_accept_prob=target_accept_prob,
              seed=seed,
              dtype=dtype,
              use_xla=use_xla,
          )
      )
  
      logging.info("Acceptance rate: %s", mcmc_results.accept_rate)
      logging.info("ESS: %s", mcmc_results.ess)
      logging.info("r_hat: %s", mcmc_results.r_hat)
  
      for (
          name,
          sample_transformation,
      ) in model.sample_transformations.items():
          transformed_chain = self.evaluate(
              tf.nest.map_structure(
                  tf.identity, sample_transformation(mcmc_results.chain)
              )
          )
  
          # tfp.mcmc.effective_sample_size only works well with lists.
          flat_transformed_chain = tf.nest.flatten(transformed_chain)
          cross_chain_dims = [1] * len(flat_transformed_chain)
          flat_ess = self.evaluate(
              tfp.mcmc.effective_sample_size(
                  flat_transformed_chain,
                  cross_chain_dims=cross_chain_dims,
                  filter_beyond_positive_pairs=True,
              )
          )
          self._z_test(
              name=name,
              sample_transformation=sample_transformation,
              transformed_samples=transformed_chain,
              num_samples=tf.nest.pack_sequence_as(
                  transformed_chain, flat_ess
              ),
              mean_fudge_atol=mean_fudge_atol,
              standard_deviation_fudge_atol=standard_deviation_fudge_atol,
              sample_dims=(0, 1),
          )
  
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 2 for calling the InferenceGymTestCase obj's underlying member method code with example usage: obj.validate_ground_truth_using_hmc(model, num_chains, num_steps, num_leapfrog_steps, step_size, target_accept_prob, seed, dtype, mean_fudge_atol, standard_deviation_fudge_atol, use_xla) without return types