code: def _ais_gets_correct_log_normalizer(
      self, init, independent_chain_ndims, use_transformed_kernel=False
  ):
      counter = collections.Counter()
  
      def proposal_log_prob(x):
          counter["proposal_calls"] += 1
          event_dims = ps.range(independent_chain_ndims, ps.rank(x))
          return tf.reduce_sum(
              normal.Normal(loc=0.0, scale=1.0).log_prob(x),
              axis=event_dims,
          )
  
      def target_log_prob(x):
          counter["target_calls"] += 1
          event_dims = ps.range(independent_chain_ndims, ps.rank(x))
          return self._log_gamma_log_prob(x, event_dims)
  
      num_steps = 200
  
      if use_transformed_kernel:
  
          def make_kernel(tlp_fn):
              return transformed_kernel.TransformedTransitionKernel(
                  inner_kernel=hmc.HamiltonianMonteCarlo(
                      target_log_prob_fn=tlp_fn,
                      step_size=0.5,
                      num_leapfrog_steps=2,
                  ),
                  bijector=identity.Identity(),
              )
  
      else:
  
          def make_kernel(tlp_fn):
              return hmc.HamiltonianMonteCarlo(
                  target_log_prob_fn=tlp_fn,
                  step_size=0.5,
                  num_leapfrog_steps=2,
              )
  
      _, ais_weights, _ = sai.sample_annealed_importance_chain(
          num_steps=num_steps,
          proposal_log_prob_fn=proposal_log_prob,
          target_log_prob_fn=target_log_prob,
          current_state=init,
          make_kernel_fn=make_kernel,
          seed=test_util.test_seed(),
      )
  
      # We have three calls because the calculation of `ais_weights` entails
      # another call to the `convex_combined_log_prob_fn`. We could refactor
      # things to avoid this, if needed (eg, b/72994218).
      if not tf.executing_eagerly():
          self.assertAllEqual(
              dict(target_calls=3, proposal_calls=3), counter
          )
  
      event_shape = tf.shape(init)[independent_chain_ndims:]
      event_size = tf.reduce_prod(event_shape)
  
      log_true_normalizer = -self._shape_param * tf.math.log(
          self._rate_param
      ) + tf.math.lgamma(self._shape_param)
      log_true_normalizer *= tf.cast(
          event_size, log_true_normalizer.dtype
      )
  
      ais_weights_size = tf.cast(
          tf.size(ais_weights), ais_weights.dtype
      )
      log_estimated_normalizer = tf.reduce_logsumexp(
          ais_weights
      ) - tf.math.log(ais_weights_size)
  
      ratio_estimate_true = tf.math.exp(
          ais_weights - log_true_normalizer
      )
      standard_error = tf.sqrt(
          tf.math.reduce_variance(ratio_estimate_true)
          / ais_weights_size
      )
  
      [
          ratio_estimate_true_,
          log_true_normalizer_,
          log_estimated_normalizer_,
          standard_error_,
          ais_weights_size_,
          event_size_,
      ] = self.evaluate(
          [
              ratio_estimate_true,
              log_true_normalizer,
              log_estimated_normalizer,
              standard_error,
              ais_weights_size,
              event_size,
          ]
      )
  
      tf1.logging.vlog(
          1,
          "        log_true_normalizer: {}\n"
          "   log_estimated_normalizer: {}\n"
          "           ais_weights_size: {}\n"
          "                 event_size: {}\n".format(
              log_true_normalizer_,
              log_estimated_normalizer_,
              ais_weights_size_,
              event_size_,
          ),
      )
      self.assertNear(
          ratio_estimate_true_.mean(), 1.0, 4.0 * standard_error_
      )
  
imports: import collections
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import sample_annealed_importance
  from tensorflow_probability.python.mcmc import transformed_kernel
label_desc: Use image node_id 3 for calling the SampleAnnealedImportanceTest obj's underlying member method code with example usage: obj._ais_gets_correct_log_normalizer(init, independent_chain_ndims, use_transformed_kernel) and returns: tf, self, transformed_kernel, hmc