def _kl_half_normal_half_normal(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b `HalfNormal`.

  Args:
    a: Instance of a `HalfNormal` distribution object.
    b: Instance of a `HalfNormal` distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_half_normal_half_normal".

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_half_normal_half_normal"):
    # Consistent with
    # http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 119
    return (tf.math.log(b.scale) - tf.math.log(a.scale) +
            (a.scale**2 - b.scale**2) / (2 * b.scale**2))
Calculate the batched KL divergence KL(a || b) with a and b `HalfNormal`.

  Args:
    a: Instance of a `HalfNormal` distribution object.
    b: Instance of a `HalfNormal` distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_half_normal_half_normal".

  Returns:
    Batchwise KL(a || b)