tensorflow_probability/python/distributions/onehot_categorical.py
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: (optional) Name to use for created operations.
      default is "kl_categorical_categorical".

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_categorical_categorical"):
    # sum(p ln(p / q))
    return tf.reduce_sum(
        input_tensor=tf.nn.softmax(a.logits) *
        (tf.nn.log_softmax(a.logits) - tf.nn.log_softmax(b.logits)),
        axis=-1)
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: (optional) Name to use for created operations.
      default is "kl_categorical_categorical".

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_categorical_categorical"):
    # sum(p ln(p / q))
    return tf.reduce_sum(
        input_tensor=tf.nn.softmax(a.logits) *
        (tf.nn.log_softmax(a.logits) - tf.nn.log_softmax(b.logits)),
        axis=-1)
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: (optional) Name to use for created operations.
      default is "kl_categorical_categorical".

  Returns:
    Batchwise KL(a || b)