class_docstr: More careful tests that sampling/preconditioning is actually working.
class_docstr_tok: ["More","careful","tests","that","sampling\/preconditioning","is","actually","working","."]
docstr: Run HMC with step_size adaptation, and return RunHMCResults.
docstr_tok: ["Run","HMC","with","step_size","adaptation",",","and","return","RunHMCResults","."]
code: def _run_hmc_with_step_size(
      self,
      target_mvn,
      precondition_scheme,
      target_accept=0.75,
      num_results=2000,
      num_adaptation_steps=20,
  ):
      """Run HMC with step_size adaptation, and return RunHMCResults."""
      assert precondition_scheme in PRECONDITION_SCHEMES
  
      dims = target_mvn.event_shape[0]
      target_cov = target_mvn.covariance()
  
      cov_linop = tf.linalg.LinearOperatorFullMatrix(
          target_cov, is_self_adjoint=True, is_positive_definite=True
      )
  
      if precondition_scheme == "no_preconditioner":
          momentum_distribution = None
          # Internal to the sampler, these scales are being used (implicitly).
          internal_scales = tf.sqrt(tf.linalg.eigvalsh(target_cov))
      elif precondition_scheme == "direct":
          momentum_distribution = mvn_linear_operator.MultivariateNormalLinearOperator(
              # The covariance of momentum is inv(covariance of position), and we
              # parameterize distributions by a square root of the covariance.
              scale=cov_linop.inverse().cholesky(),
          )
          # Internal to the sampler, these scales are being used (implicitly).
          internal_scales = tf.ones(dims)
      elif precondition_scheme == "precision_factor":
          momentum_distribution = mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
              # The precision of momentum is the covariance of position.
              # The "factor" is the cholesky factor.
              precision_factor=cov_linop.cholesky(),
          )
          # Internal to the sampler, these scales are being used (implicitly).
          internal_scales = tf.ones(dims)
      elif precondition_scheme == "sqrtm":
          if JAX_MODE:
              self.skipTest("`sqrtm` is not yet implemented in JAX.")
          momentum_distribution = mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
              # The symmetric square root is a perfectly valid "factor".
              precision_factor=tf.linalg.LinearOperatorFullMatrix(
                  tf.linalg.sqrtm(target_cov)
              ),
          )
          # Internal to the sampler, these scales are being used (implicitly).
          internal_scales = tf.ones(dims)
      elif precondition_scheme == "scale":
          momentum_distribution = mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
              # Nothing wrong with using "scale", since the scale should be the
              # same as cov_linop.cholesky().
              precision_factor=target_mvn.scale,
          )
          # Internal to the sampler, these scales are being used (implicitly).
          internal_scales = tf.ones(dims)
      else:
          raise RuntimeError(
              "Unhandled precondition_scheme: {}".format(
                  precondition_scheme
              )
          )
      momentum_distribution = _make_composite_tensor(
          momentum_distribution
      )
  
      # Asyptotic step size, assuming P[accept] = target_accept.
      expected_step = self._calculate_asymptotic_step_size(
          scales=internal_scales,
          prob_accept=target_accept,
      )
  
      # Initialize step size to something close to the expected required step
      # size. This helps reduce the need for a long burn-in. Don't use the
      # expected step size exactly, since that would be cheating.
      initial_step_size = expected_step / 2.345
  
      # Set num_leapfrog_steps so that we get decent ESS.
      max_internal_scale = tf.reduce_max(internal_scales)
      num_leapfrog_steps = tf.minimum(
          tf.cast(
              tf.math.ceil(1.5 * max_internal_scale / expected_step),
              dtype=tf.int32,
          ),
          30,
      )
  
      hmc_kernel = dassa.DualAveragingStepSizeAdaptation(
          phmc.PreconditionedHamiltonianMonteCarlo(
              target_log_prob_fn=target_mvn.log_prob,
              momentum_distribution=momentum_distribution,
              step_size=initial_step_size,
              num_leapfrog_steps=num_leapfrog_steps,
          ),
          num_adaptation_steps=num_adaptation_steps,
          target_accept_prob=target_accept,
      )
  
      def trace_fn(_, pkr):
          results = pkr.inner_results
          return {
              "accept_prob": tf.exp(
                  tf.minimum(0.0, results.log_accept_ratio)
              ),
              "step_size": results.accepted_results.step_size,
          }
  
      @tf.function
      def do_run_run_run():
          """Do a run, return RunHMCResults."""
          states, trace = sample.sample_chain(
              num_results,
              current_state=tf.identity(
                  target_mvn.sample(seed=test_util.test_seed())
              ),
              kernel=hmc_kernel,
              num_burnin_steps=num_adaptation_steps,
              seed=test_util.test_seed(),
              trace_fn=trace_fn,
          )
  
          # If we had some number of chain dimensions, we would change sample_axis.
          sample_axis = 0
  
          sample_cov = sample_stats.covariance(
              states, sample_axis=sample_axis
          )
          max_variance = tf.reduce_max(tf.linalg.diag_part(sample_cov))
          max_stddev = tf.sqrt(max_variance)
          min_ess = tf.reduce_min(
              diagnostic.effective_sample_size(states)
          )
          mean_accept_prob = tf.reduce_mean(trace["accept_prob"])
  
          # Asymptotic step size given that P[accept] = mean_accept_prob.
          asymptotic_step_size = self._calculate_asymptotic_step_size(
              scales=internal_scales,
              prob_accept=mean_accept_prob,
          )
  
          return RunHMCResults(
              draws=states,
              step_size=trace["step_size"],
              final_step_size=trace["step_size"][-1],
              asymptotic_step_size=asymptotic_step_size,
              accept_prob=trace["accept_prob"],
              mean_accept_prob=mean_accept_prob,
              min_ess=tf.reduce_min(
                  diagnostic.effective_sample_size(states)
              ),
              sample_mean=tf.reduce_mean(states, axis=sample_axis),
              sample_cov=sample_cov,
              sample_var=tf.linalg.diag_part(sample_cov),
              # Standard error in variance estimation is related to standard
              # deviation of variance estimates. For a Normal, this is just Sqrt(2)
              # times variance divided by sqrt sample size (or so my old notes say).
              # So a relative tolerance is useful.
              # Add in a factor of 5 as a buffer.
              var_rtol=5 * tf.sqrt(2.0) / tf.sqrt(min_ess),
              # For covariance matrix estimates, there can be terms that have
              # expectation = 0 (e.g. off diagonal entries). So the above doesn't
              # hold. So use an atol.
              cov_atol=5 * max_variance / tf.sqrt(min_ess),
              # Standard error in mean estimation is stddev divided by sqrt
              # sample size. This is an absolute tolerance.
              # Add in a factor of 5 as a buffer.
              mean_atol=5 * max_stddev / tf.sqrt(min_ess),
          )
  
      # Evaluate now, to ensure that states/accept_prob/etc... all match up with
      # the same graph evaluation. This is a gotcha about TFP MCMC in graph mode.
      return self.evaluate(do_run_run_run())
  
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 2 for calling the PreconditionedHMCCorrectnessTest obj's underlying member method code with example usage: obj._run_hmc_with_step_size(target_mvn, precondition_scheme, target_accept, num_results, num_adaptation_steps) and returns: self, dict, RunHMCResults