def _kl_independent(a, b, name="kl_independent"):
  """Batched KL divergence `KL(a || b)` for Independent distributions.

  We can leverage the fact that
  ```
  KL(Independent(a) || Independent(b)) = sum(KL(a || b))
  ```
  where the sum is over the `reinterpreted_batch_ndims`.

  Args:
    a: Instance of `Independent`.
    b: Instance of `Independent`.
    name: (optional) name to use for created ops. Default "kl_independent".

  Returns:
    Batchwise `KL(a || b)`.

  Raises:
    ValueError: If the event space for `a` and `b`, or their underlying
      distributions don't match.
  """
  p = a.distribution
  q = b.distribution

  # The KL between any two (non)-batched distributions is a scalar.
  # Given that the KL between two factored distributions is the sum, i.e.
  # KL(p1(x)p2(y) || q1(x)q2(y)) = KL(p1 || q1) + KL(q1 || q2), we compute
  # KL(p || q) and do a `reduce_sum` on the reinterpreted batch dimensions.
  if (tensorshape_util.is_fully_defined(a.event_shape) and
      tensorshape_util.is_fully_defined(b.event_shape)):
    if a.event_shape == b.event_shape:
      if p.event_shape == q.event_shape:
        num_reduce_dims = (tensorshape_util.rank(a.event_shape) -
                           tensorshape_util.rank(p.event_shape))
        reduce_dims = [-i - 1 for i in range(0, num_reduce_dims)]

        return tf.reduce_sum(
            input_tensor=kullback_leibler.kl_divergence(p, q, name=name),
            axis=reduce_dims)
      else:
        raise NotImplementedError("KL between Independents with different "
                                  "event shapes not supported.")
    else:
      raise ValueError("Event shapes do not match.")
  else:
    with tf.control_dependencies(
        [
            assert_util.assert_equal(a.event_shape_tensor(),
                                     b.event_shape_tensor()),
            assert_util.assert_equal(p.event_shape_tensor(),
                                     q.event_shape_tensor())
        ]):
      num_reduce_dims = (
          prefer_static.rank_from_shape(
              a.event_shape_tensor, a.event_shape) -
          prefer_static.rank_from_shape(
              p.event_shape_tensor, a.event_shape))
      reduce_dims = prefer_static.range(-num_reduce_dims - 1, -1, 1)
      return tf.reduce_sum(
          input_tensor=kullback_leibler.kl_divergence(p, q, name=name),
          axis=reduce_dims)
def _kl_independent(a, b, name="kl_independent"):
  """Batched KL divergence `KL(a || b)` for Independent distributions.

  We can leverage the fact that
  ```
  KL(Independent(a) || Independent(b)) = sum(KL(a || b))
  ```
  where the sum is over the `reinterpreted_batch_ndims`.

  Args:
    a: Instance of `Independent`.
    b: Instance of `Independent`.
    name: (optional) name to use for created ops. Default "kl_independent".

  Returns:
    Batchwise `KL(a || b)`.

  Raises:
    ValueError: If the event space for `a` and `b`, or their underlying
      distributions don't match.
  """
  p = a.distribution
  q = b.distribution

  # The KL between any two (non)-batched distributions is a scalar.
  # Given that the KL between two factored distributions is the sum, i.e.
  # KL(p1(x)p2(y) || q1(x)q2(y)) = KL(p1 || q1) + KL(q1 || q2), we compute
  # KL(p || q) and do a `reduce_sum` on the reinterpreted batch dimensions.
  if (tensorshape_util.is_fully_defined(a.event_shape) and
      tensorshape_util.is_fully_defined(b.event_shape)):
    if a.event_shape == b.event_shape:
      if p.event_shape == q.event_shape:
        num_reduce_dims = (tensorshape_util.rank(a.event_shape) -
                           tensorshape_util.rank(p.event_shape))
        reduce_dims = [-i - 1 for i in range(0, num_reduce_dims)]

        return tf.reduce_sum(
            input_tensor=kullback_leibler.kl_divergence(p, q, name=name),
            axis=reduce_dims)
      else:
        raise NotImplementedError("KL between Independents with different "
                                  "event shapes not supported.")
    else:
      raise ValueError("Event shapes do not match.")
  else:
    with tf.control_dependencies(
        [
            assert_util.assert_equal(a.event_shape_tensor(),
                                     b.event_shape_tensor()),
            assert_util.assert_equal(p.event_shape_tensor(),
                                     q.event_shape_tensor())
        ]):
      num_reduce_dims = (
          prefer_static.rank_from_shape(
              a.event_shape_tensor, a.event_shape) -
          prefer_static.rank_from_shape(
              p.event_shape_tensor, a.event_shape))
      reduce_dims = prefer_static.range(-num_reduce_dims - 1, -1, 1)
      return tf.reduce_sum(
          input_tensor=kullback_leibler.kl_divergence(p, q, name=name),
          axis=reduce_dims)
def _kl_independent ( a , b , name = "kl_independent" ) : p = a . distribution q = b . distribution # The KL between any two (non)-batched distributions is a scalar. # Given that the KL between two factored distributions is the sum, i.e. # KL(p1(x)p2(y) || q1(x)q2(y)) = KL(p1 || q1) + KL(q1 || q2), we compute # KL(p || q) and do a `reduce_sum` on the reinterpreted batch dimensions. if ( tensorshape_util . is_fully_defined ( a . event_shape ) and tensorshape_util . is_fully_defined ( b . event_shape ) ) : if a . event_shape == b . event_shape : if p . event_shape == q . event_shape : num_reduce_dims = ( tensorshape_util . rank ( a . event_shape ) - tensorshape_util . rank ( p . event_shape ) ) reduce_dims = [ - i - 1 for i in range ( 0 , num_reduce_dims ) ] return tf . reduce_sum ( input_tensor = kullback_leibler . kl_divergence ( p , q , name = name ) , axis = reduce_dims ) else : raise NotImplementedError ( "KL between Independents with different " "event shapes not supported." ) else : raise ValueError ( "Event shapes do not match." ) else : with tf . control_dependencies ( [ assert_util . assert_equal ( a . event_shape_tensor ( ) , b . event_shape_tensor ( ) ) , assert_util . assert_equal ( p . event_shape_tensor ( ) , q . event_shape_tensor ( ) ) ] ) : num_reduce_dims = ( prefer_static . rank_from_shape ( a . event_shape_tensor , a . event_shape ) - prefer_static . rank_from_shape ( p . event_shape_tensor , a . event_shape ) ) reduce_dims = prefer_static . range ( - num_reduce_dims - 1 , - 1 , 1 ) return tf . reduce_sum ( input_tensor = kullback_leibler . kl_divergence ( p , q , name = name ) , axis = reduce_dims )
Batched KL divergence `KL(a || b)` for Independent distributions.

  We can leverage the fact that
  ```
  KL(Independent(a) || Independent(b)) = sum(KL(a || b))
  ```
  where the sum is over the `reinterpreted_batch_ndims`.

  Args:
    a: Instance of `Independent`.
    b: Instance of `Independent`.
    name: (optional) name to use for created ops. Default "kl_independent".

  Returns:
    Batchwise `KL(a || b)`.

  Raises:
    ValueError: If the event space for `a` and `b`, or their underlying
      distributions don't match.
Batched KL divergence KL ( a || b ) for Independent distributions .