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

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

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_bernoulli_bernoulli"):
    delta_probs0 = tf.nn.softplus(-b.logits) - tf.nn.softplus(-a.logits)
    delta_probs1 = tf.nn.softplus(b.logits) - tf.nn.softplus(a.logits)
    return (tf.sigmoid(a.logits) * delta_probs0
            + tf.sigmoid(-a.logits) * delta_probs1)
Calculate the batched KL divergence KL(a || b) with a and b Bernoulli.

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

  Returns:
    Batchwise KL(a || b)