code: def testGumbelGumbelKL(self):
      a_loc = np.arange(-2.0, 3.0, 1.0)
      a_scale = np.arange(0.5, 2.5, 0.5)
      b_loc = 2 * np.arange(-2.0, 3.0, 1.0)
      b_scale = np.arange(0.5, 2.5, 0.5)
  
      # This reshape is intended to expand the number of test cases.
      a_loc = a_loc.reshape((len(a_loc), 1, 1, 1))
      a_scale = a_scale.reshape((1, len(a_scale), 1, 1))
      b_loc = b_loc.reshape((1, 1, len(b_loc), 1))
      b_scale = b_scale.reshape((1, 1, 1, len(b_scale)))
  
      a = gumbel.Gumbel(loc=a_loc, scale=a_scale, validate_args=True)
      b = gumbel.Gumbel(loc=b_loc, scale=b_scale, validate_args=True)
  
      true_kl = (
          np.log(b_scale)
          - np.log(a_scale)
          + np.euler_gamma * (a_scale / b_scale - 1.0)
          + np.expm1(
              (b_loc - a_loc) / b_scale
              + np.vectorize(np.math.lgamma)(a_scale / b_scale + 1.0)
          )
          + (a_loc - b_loc) / b_scale
      )
  
      kl = kullback_leibler.kl_divergence(a, b)
  
      x = a.sample(int(1e5), seed=test_util.test_seed())
      kl_sample = tf.reduce_mean(a.log_prob(x) - b.log_prob(x), axis=0)
  
      # As noted in the Gumbel-Gumbel KL divergence implementation, there is an
      # error in the reference paper we use to implement our divergence. This
      # error is a missing summand, (a.loc - b.loc) / b.scale. To ensure that we
      # are adequately testing this difference in the below tests, we compute the
      # relative error between kl_sample_ and kl_ and check that it is "much less"
      # than this missing summand.
      summand = (a_loc - b_loc) / b_scale
      relative_error = tf.abs(kl - kl_sample) / tf.minimum(
          tf.abs(kl), tf.abs(kl_sample)
      )
      exists_missing_summand_test = tf.reduce_any(
          summand > 2 * relative_error
      )
      exists_missing_summand_test_ = self.evaluate(
          exists_missing_summand_test
      )
      self.assertTrue(
          exists_missing_summand_test_,
          msg=(
              "No test case exists where (a.loc - b.loc) / b.scale "
              "is much less than the relative error between kl as "
              "computed in closed form, and kl as computed by "
              "sampling. Failing to include such a test case makes "
              "it difficult to detect regressions where this "
              "summand (which is missing in our reference paper) "
              "is omitted."
          ),
      )
  
      kl_, kl_sample_ = self.evaluate([kl, kl_sample])
      self.assertAllClose(true_kl, kl_, atol=1e-15, rtol=1e-12)
      self.assertAllClose(true_kl, kl_sample_, atol=0.0, rtol=1e-1)
  
      zero_kl = kullback_leibler.kl_divergence(a, a)
      true_zero_kl_, zero_kl_ = self.evaluate(
          [tf.zeros_like(zero_kl), zero_kl]
      )
      # TODO(b/157595661): Investigate why this tolerance is needed.
      self.assertAllClose(true_zero_kl_, zero_kl_, atol=1e-15)
  
imports: import numpy
  from scipy import stats
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import gumbel
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 15 for calling the _GumbelTest obj's underlying member method code with example usage: obj.testGumbelGumbelKL() without return types