tensorflow_probability/python/distributions/gumbel.py
def _kl_gumbel_gumbel(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b Gumbel.

  Args:
    a: instance of a Gumbel distribution object.
    b: instance of a Gumbel distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_gumbel_gumbel".

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_gumbel_gumbel"):
    # Consistent with
    # http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 64
    # The paper uses beta to refer to scale and mu to refer to loc.
    # There is actually an error in the solution as printed; this is based on
    # the second-to-last step of the derivation. The value as printed would be
    # off by (a.loc - b.loc) / b.scale.
    return (tf.math.log(b.scale) - tf.math.log(a.scale) + np.euler_gamma *
            (a.scale / b.scale - 1.) +
            tf.math.expm1((b.loc - a.loc) / b.scale +
                          tf.math.lgamma(a.scale / b.scale + 1.)) +
            (a.loc - b.loc) / b.scale)
def _kl_gumbel_gumbel(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b Gumbel.

  Args:
    a: instance of a Gumbel distribution object.
    b: instance of a Gumbel distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_gumbel_gumbel".

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_gumbel_gumbel"):
    # Consistent with
    # http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 64
    # The paper uses beta to refer to scale and mu to refer to loc.
    # There is actually an error in the solution as printed; this is based on
    # the second-to-last step of the derivation. The value as printed would be
    # off by (a.loc - b.loc) / b.scale.
    return (tf.math.log(b.scale) - tf.math.log(a.scale) + np.euler_gamma *
            (a.scale / b.scale - 1.) +
            tf.math.expm1((b.loc - a.loc) / b.scale +
                          tf.math.lgamma(a.scale / b.scale + 1.)) +
            (a.loc - b.loc) / b.scale)
Calculate the batched KL divergence KL(a || b) with a and b Gumbel.

  Args:
    a: instance of a Gumbel distribution object.
    b: instance of a Gumbel distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_gumbel_gumbel".

  Returns:
    Batchwise KL(a || b)