docstr: Make dictionary of results and do some standard checks.
code: def checkAndMakeResultsForTestingUntemperedLogProbFn(
      self,
      likelihood_variance,
      prior_variance,
      inverse_temperatures,
      num_results,
      expected_min_conditional_swap_prob=None,
      expected_max_conditional_swap_prob=None,
      tol_multiplier=10,
  ):
      """Make dictionary of results and do some standard checks."""
      # With an untempered log prob fn we can use beta = 0. This unfortunately
      # requires a different setup in terms of step size than the usual tests, and
      # we expect slightly different results. So we made our own special test.
      inverse_temperatures = tf.convert_to_tensor(inverse_temperatures)
      num_replica = inverse_temperatures.shape[0]
      dtype = inverse_temperatures.dtype
  
      is_list_like = isinstance(likelihood_variance, (list, tuple))
      # If is_list_like, we will concatenate tensors along the way.
      # The log prob fns we build will do this concatenation to inputs
      # The "theoretical" outputs will be concatenated.
      # It's up to the user to split back into "parts" if required.
  
      if is_list_like:
          assert isinstance(prior_variance, (list, tuple))
          assert [len(v.shape) == 1 for v in likelihood_variance]
          assert [len(v.shape) == 1 for v in prior_variance]
          n_events = [v.shape[-1] for v in likelihood_variance]
          likelihood_variance = tf.concat(likelihood_variance, axis=-1)
          prior_variance = tf.concat(prior_variance, axis=-1)
          current_state = [tf.zeros((num_replica, n)) for n in n_events]
      else:
          assert len(likelihood_variance.shape) == 1
          assert len(prior_variance.shape) == 1
          n_events = likelihood_variance.shape[-1]
          likelihood_variance = tf.convert_to_tensor(
              likelihood_variance, dtype=dtype
          )
          prior_variance = tf.convert_to_tensor(
              prior_variance, dtype=dtype
          )
          current_state = tf.zeros((num_replica, n_events))
  
      likelihood = mvn_diag.MultivariateNormalDiag(
          scale_diag=tf.sqrt(likelihood_variance)
      )
      prior = mvn_diag.MultivariateNormalDiag(
          scale_diag=tf.sqrt(prior_variance)
      )
  
      def likelihood_log_prob(*x):
          if is_list_like:
              x = tf.concat(x, axis=-1)
          else:
              x = x[0]
          return likelihood.log_prob(x)
  
      def prior_log_prob(*x):
          if is_list_like:
              x = tf.concat(x, axis=-1)
          else:
              x = x[0]
          return prior.log_prob(x)
  
      # Precision is inverse of covariance. The (untempered) prior and
      # (tempered) likelihood precisions add.
      replica_precision = tf.linalg.LinearOperatorDiag(
          tf.stack(
              [
                  beta / likelihood_variance + 1.0 / prior_variance
                  for beta in tf.unstack(inverse_temperatures)
              ]
          )
      )
      replica_covariance = replica_precision.inverse()
  
      replica_smallest_scale_length = tf.reduce_min(
          tf.sqrt(replica_covariance.diag_part()), axis=-1
      )
      step_size = 0.41234 * replica_smallest_scale_length[:, tf.newaxis]
      num_leapfrog_steps = tf.cast(
          tf.math.ceil(1.75 / tf.reduce_min(step_size)), tf.int32
      )
  
      def make_kernel_fn(target_log_prob_fn):
          return hmc.HamiltonianMonteCarlo(
              target_log_prob_fn=target_log_prob_fn,
              step_size=step_size,
              num_leapfrog_steps=num_leapfrog_steps,
          )
  
      remc = replica_exchange_mc.ReplicaExchangeMC(
          target_log_prob_fn=None,
          untempered_log_prob_fn=prior_log_prob,
          tempered_log_prob_fn=likelihood_log_prob,
          inverse_temperatures=inverse_temperatures,
          state_includes_replicas=True,
          make_kernel_fn=make_kernel_fn,
          swap_proposal_fn=replica_exchange_mc.even_odd_swap_proposal_fn(
              1.0
          ),
      )
  
      def trace_fn(state, results):  # pylint: disable=unused-argument
          return {
              "replica_log_accept_ratio": unnest.get_innermost(
                  results.post_swap_replica_results, "log_accept_ratio"
              ),
              "is_swap_accepted_adjacent": results.is_swap_accepted_adjacent,
              "is_swap_proposed_adjacent": results.is_swap_proposed_adjacent,
          }
  
      replica_states, trace = sample.sample_chain(
          num_results=num_results,
          # Start at one of the modes, in order to make mode jumping necessary
          # if we want to pass test.
          current_state=current_state,
          kernel=remc,
          num_burnin_steps=50,
          trace_fn=trace_fn,
          seed=test_util.test_seed(),
      )
  
      # Make sure replicas are swapping. If they are not, then the individual
      # replicas will converge even if the swap probabilities are incorrect.
      conditional_swap_prob = tf.reduce_sum(
          tf.cast(trace["is_swap_accepted_adjacent"], dtype), axis=0
      ) / tf.reduce_sum(
          tf.cast(trace["is_swap_proposed_adjacent"], dtype), axis=0
      )
      replica_prob_accept = tf.reduce_mean(
          tf.exp(tf.minimum(trace["replica_log_accept_ratio"], 0.0)),
          axis=0,
      )
  
      if is_list_like:
          replica_states = tf.concat(replica_states, axis=-1)
  
      results = self.evaluate(
          {
              "replica_states": replica_states,
              "replica_prob_accept": replica_prob_accept,
              "conditional_swap_prob": conditional_swap_prob,
              "theoretical_replica_mean": tf.zeros_like(
                  replica_covariance.diag_part()
              ),
              "theoretical_replica_variance": replica_covariance.diag_part(),
              "sample_replica_variance": sample_stats.variance(
                  replica_states, sample_axis=0
              ),
              "sample_replica_mean": tf.reduce_mean(
                  replica_states, axis=0
              ),
              "replica_ess": effective_sample_size(replica_states),
          }
      )
  
      results["min_ess"] = np.min(results["replica_ess"])
  
      # All replicas should be mixing well, individually.
      self.assertGreater(
          results["min_ess"],
          num_results / 10,
          msg="Bad sampling found!",
      )
      self.assertAllGreater(results["replica_prob_accept"], 0.6)
  
      # If replicas are not swapping properly, other results may look nice
      # even if swap probabilities are incorrect.
      if expected_min_conditional_swap_prob is not None:
          self.assertAllGreater(
              results["conditional_swap_prob"],
              expected_min_conditional_swap_prob,
          )
      if expected_max_conditional_swap_prob is not None:
          self.assertAllLess(
              results["conditional_swap_prob"],
              expected_max_conditional_swap_prob,
          )
  
      mean_atol = tol_multiplier / np.sqrt(results["min_ess"])
      cov_rtol = (
          tol_multiplier * np.sqrt(2) / np.sqrt(results["min_ess"])
      )
  
      self.assertAllClose(
          results["sample_replica_variance"],
          results["theoretical_replica_variance"],
          rtol=cov_rtol,
      )
  
      self.assertAllClose(
          results["sample_replica_mean"],
          results["theoretical_replica_mean"],
          atol=mean_atol,
      )
  
      return results
  
imports: from absl import logging
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_combinations
  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 hmc
  from tensorflow_probability.python.mcmc import random_walk_metropolis
  from tensorflow_probability.python.mcmc import replica_exchange_mc
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 18 for calling the REMCTest obj's underlying member method code with example usage: obj.checkAndMakeResultsForTestingUntemperedLogProbFn(likelihood_variance, prior_variance, inverse_temperatures, num_results, expected_min_conditional_swap_prob, expected_max_conditional_swap_prob, tol_multiplier) and returns: results, likelihood, prior, hmc, dict