code: def _kl_categorical_categorical(a, b, name=None):
      """Calculate the batched KL divergence KL(a || b) with a, b OneHotCategorical.
  
      Args:
        a: instance of a OneHotCategorical distribution object.
        b: instance of a OneHotCategorical distribution object.
        name: Python `str` name to use for created operations.
          Default value: `None` (i.e., `'kl_categorical_categorical'`).
  
      Returns:
        Batchwise KL(a || b)
      """
      with tf.name_scope(name or "kl_categorical_categorical"):
          # pylint: disable=protected-access
          a_logits = a._logits_parameter_no_checks()
          # pylint: disable=protected-access
          b_logits = b._logits_parameter_no_checks()
          # sum(p ln(p / q))
          return tf.reduce_sum(
              (
                  tf.math.softmax(a_logits)
                  * (
                      tf.math.log_softmax(a_logits)
                      - tf.math.log_softmax(b_logits)
                  )
              ),
              axis=-1,
          )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softmax_centered
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: _kl_categorical_categorical(a, b, name) and returns: tf