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

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

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_chi_chi"):
    # Consistent with
    # https://mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 118
    # The paper introduces an additional scaling parameter; setting that
    # parameter to 1 and simplifying yields the expression we use here.
    return (0.5 * tf.math.digamma(0.5 * a.df) * (a.df - b.df) +
            tf.math.lgamma(0.5 * b.df) - tf.math.lgamma(0.5 * a.df))
def _kl_chi_chi(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b Chi.

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

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_chi_chi"):
    # Consistent with
    # https://mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 118
    # The paper introduces an additional scaling parameter; setting that
    # parameter to 1 and simplifying yields the expression we use here.
    return (0.5 * tf.math.digamma(0.5 * a.df) * (a.df - b.df) +
            tf.math.lgamma(0.5 * b.df) - tf.math.lgamma(0.5 * a.df))
Calculate the batched KL divergence KL(a || b) with a and b Chi.

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

  Returns:
    Batchwise KL(a || b)