class_docstr: More careful tests that sampling/preconditioning is actually working.
class_docstr_tok: ["More","careful","tests","that","sampling\/preconditioning","is","actually","working","."]
docstr: Test that step size adaptation finds the theoretical optimal step size.
  
  See _caclulate_expected_step_size for formula details, but roughly, for a
  high dimensional Gaussian posterior, we can calculate the approximate step
  size to achieve a given target accept rate. For such a posterior,
  `PreconditionedHMC` mimics the dynamics of sampling from an isotropic
  standard normal distribution, and so should adapt to the step size where
  the scales are all ones.
  
  In the example below, `expected_step` is around 0.00002, so there is
  significantly different behavior when conditioning.
  
  Args:
    target_mvn: Multivariate normal instance to sample from.
    num_results: Number of samples to collect (post burn-in).
    precondition_scheme: String telling how to do preconditioning.
      Should be in PRECONDITION_SCHEMES.
    check_step_size_asymptotics: Boolean telling whether to check that the
      step size and P[accept] match up with expected values. This checks
      that the "internal/implicit" sampling distribution is as expected. E.g.
      when preconditioning, we expect the internal distribution to be a
      standard Normal. When not preconditioning we expect it to be the target.
    asymptotic_step_size_rtol: rtol for the asymptotic step size test.
      The "nastier" spectra (with a small number of tiny eigenvalues) often
      require larger tolerance.  About 10% rtol is what we can expect.
      20% is the default for safety.  When a "bad preconditioner" is used,
      these two are off by 100% or more (but no guarantee, since luck may
      prevail).
  
  Returns:
    RunHMCResults
docstr_tok: ["Test","that","step","size","adaptation","finds","the","theoretical","optimal","step","size",".","See","_caclulate_expected_step_size","for","formula","details",",","but","roughly",",","for","a","high","dimensional","Gaussian","posterior",",","we","can","calculate","the","approximate","step","size","to","achieve","a","given","target","accept","rate",".","For","such","a","posterior",",","`","PreconditionedHMC","`","mimics","the","dynamics","of","sampling","from","an","isotropic","standard","normal","distribution",",","and","so","should","adapt","to","the","step","size","where","the","scales","are","all","ones",".","In","the","example","below",",","`","expected_step","`","is","around","0.00002",",","so","there","is","significantly","different","behavior","when","conditioning",".","Args",":","target_mvn",":","Multivariate","normal","instance","to","sample","from",".","num_results",":","Number","of","samples","to","collect","(","post","burn-in",")",".","precondition_scheme",":","String","telling","how","to","do","preconditioning",".","Should","be","in","PRECONDITION_SCHEMES",".","check_step_size_asymptotics",":","Boolean","telling","whether","to","check","that","the","step","size","and","P","[","accept","]","match","up","with","expected","values",".","This","checks","that","the","``","internal\/implicit","''","sampling","distribution","is","as","expected",".","E.g",".","when","preconditioning",",","we","expect","the","internal","distribution","to","be","a","standard","Normal",".","When","not","preconditioning","we","expect","it","to","be","the","target",".","asymptotic_step_size_rtol",":","rtol","for","the","asymptotic","step","size","test",".","The","``","nastier","''","spectra","(","with","a","small","number","of","tiny","eigenvalues",")","often","require","larger","tolerance",".","About","10","%","rtol","is","what","we","can","expect",".","20","%","is","the","default","for","safety",".","When","a","``","bad","preconditioner","''","is","used",",","these","two","are","off","by","100","%","or","more","(","but","no","guarantee",",","since","luck","may","prevail",")",".","Returns",":","RunHMCResults"]
code: def _check_correctness_of_moments_and_preconditioning(
      self,
      target_mvn,
      num_results,
      precondition_scheme,
      check_step_size_asymptotics=True,
      asymptotic_step_size_rtol=0.2,
  ):
      """Test that step size adaptation finds the theoretical optimal step size.
  
      See _caclulate_expected_step_size for formula details, but roughly, for a
      high dimensional Gaussian posterior, we can calculate the approximate step
      size to achieve a given target accept rate. For such a posterior,
      `PreconditionedHMC` mimics the dynamics of sampling from an isotropic
      standard normal distribution, and so should adapt to the step size where
      the scales are all ones.
  
      In the example below, `expected_step` is around 0.00002, so there is
      significantly different behavior when conditioning.
  
      Args:
        target_mvn: Multivariate normal instance to sample from.
        num_results: Number of samples to collect (post burn-in).
        precondition_scheme: String telling how to do preconditioning.
          Should be in PRECONDITION_SCHEMES.
        check_step_size_asymptotics: Boolean telling whether to check that the
          step size and P[accept] match up with expected values. This checks
          that the "internal/implicit" sampling distribution is as expected. E.g.
          when preconditioning, we expect the internal distribution to be a
          standard Normal. When not preconditioning we expect it to be the target.
        asymptotic_step_size_rtol: rtol for the asymptotic step size test.
          The "nastier" spectra (with a small number of tiny eigenvalues) often
          require larger tolerance.  About 10% rtol is what we can expect.
          20% is the default for safety.  When a "bad preconditioner" is used,
          these two are off by 100% or more (but no guarantee, since luck may
          prevail).
  
      Returns:
        RunHMCResults
      """
      results = self._run_hmc_with_step_size(
          target_mvn, precondition_scheme=precondition_scheme
      )
  
      if check_step_size_asymptotics:
          self.assertAllClose(
              results.final_step_size,
              results.asymptotic_step_size,
              rtol=asymptotic_step_size_rtol,
          )
  
      self.assertAllClose(
          results.sample_mean, target_mvn.mean(), atol=results.mean_atol
      )
      self.assertAllClose(
          results.sample_var,
          target_mvn.variance(),
          rtol=results.var_rtol,
      )
      self.assertAllClose(
          results.sample_cov,
          target_mvn.covariance(),
          atol=results.cov_atol,
      )
  
      return results
  
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 3 for calling the PreconditionedHMCCorrectnessTest obj's underlying member method code with example usage: obj._check_correctness_of_moments_and_preconditioning(target_mvn, num_results, precondition_scheme, check_step_size_asymptotics, asymptotic_step_size_rtol) and returns: results