code: def testCategoricalCategoricalKLExtremeGrad(self, a_logits, b_logits):
      def get_kl_fn(a_logits, b_logits):
          a = categorical.Categorical(logits=a_logits)
          b = categorical.Categorical(logits=b_logits)
          return a.kl_divergence(b)
  
      kl, grads = gradient.value_and_gradient(
          get_kl_fn, a_logits, b_logits
      )
      # We get the expected answer when clipping the -inf.
      expected_kl, expected_grads = gradient.value_and_gradient(
          get_kl_fn,
          np.maximum(-10000, a_logits),
          np.maximum(-10000, b_logits),
      )
      self.assertAllClose(kl, expected_kl)
      self.assertAllClose(grads, expected_grads)
  
imports: 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 categorical
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 32 for calling the CategoricalTest obj's underlying member method code with example usage: obj.testCategoricalCategoricalKLExtremeGrad(a_logits, b_logits) and returns: a