docstr: Test that a diagonal multivariate normal can be effectively sampled from.
  
  Note that the effective sample size is expected to be exactly 100: this is
  because the step size is tuned well enough that a single HMC step takes
  a point to nearly the antipodal point, which causes a negative lag 1
  autocorrelation, and the effective sample size calculation cuts off when
  the autocorrelation drops below zero.
docstr_tok: ["Test","that","a","diagonal","multivariate","normal","can","be","effectively","sampled","from",".","Note","that","the","effective","sample","size","is","expected","to","be","exactly","100",":","this","is","because","the","step","size","is","tuned","well","enough","that","a","single","HMC","step","takes","a","point","to","nearly","the","antipodal","point",",","which","causes","a","negative","lag","1","autocorrelation",",","and","the","effective","sample","size","calculation","cuts","off","when","the","autocorrelation","drops","below","zero","."]
code: def test_diag(self):
      """Test that a diagonal multivariate normal can be effectively sampled from.
  
      Note that the effective sample size is expected to be exactly 100: this is
      because the step size is tuned well enough that a single HMC step takes
      a point to nearly the antipodal point, which causes a negative lag 1
      autocorrelation, and the effective sample size calculation cuts off when
      the autocorrelation drops below zero.
      """
      mvn = mvn_diag.MultivariateNormalDiag(
          loc=[1.0, 2.0, 3.0], scale_diag=[0.1, 1.0, 10.0]
      )
  
      if self.use_default_momentum_distribution:
          momentum_distribution = None
          step_size = 0.1
      else:
          momentum_distribution = (
              mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
                  precision_factor=mvn.scale,
              )
          )
          step_size = 0.3
      hmc_kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=mvn.log_prob,
          momentum_distribution=momentum_distribution,
          step_size=step_size,
          num_leapfrog_steps=10,
      )
      draws = sample.sample_chain(
          110,
          tf.zeros(3),
          kernel=hmc_kernel,
          seed=test_util.test_seed(),
          trace_fn=None,
      )
      ess = diagnostic.effective_sample_size(
          draws[-100:],
          filter_threshold=0,
          filter_beyond_positive_pairs=False,
      )
  
      if not self.use_default_momentum_distribution:
          self.assertAllClose(ess, tf.fill([3], 100.0))
      else:
          self.assertLess(self.evaluate(tf.reduce_min(ess)), 100.0)
  
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 4 for calling the _PreconditionedHMCTest obj's underlying member method code with example usage: obj.test_diag() without return types