def _kl_sample(a, b, name='kl_sample'):
  """Batched KL divergence `KL(a || b)` for Sample distributions.

  We can leverage the fact that:

  ```
  KL(Sample(a) || Sample(b)) = sum(KL(a || b))
  ```

  where the sum is over the `sample_shape` dims.

  Args:
    a: Instance of `Sample` distribution.
    b: Instance of `Sample` distribution.
    name: (optional) name to use for created ops.
      Default value: `"kl_sample"`'.

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

  Raises:
    ValueError: If the `sample_shape` of `a` and `b` don't match.
  """
  assertions = []
  a_ss = tf.get_static_value(a.sample_shape)
  b_ss = tf.get_static_value(b.sample_shape)
  msg = '`a.sample_shape` must be identical to `b.sample_shape`.'
  if a_ss is not None and b_ss is not None:
    if not np.array_equal(a_ss, b_ss):
      raise ValueError(msg)
  elif a.validate_args or b.validate_args:
    assertions.append(assert_util.assert_equal(
        a.sample_shape, b.sample_shape, message=msg))
  with tf.control_dependencies(assertions):
    kl = kullback_leibler.kl_divergence(
        a.distribution, b.distribution, name=name)
    n = prefer_static.reduce_prod(a.sample_shape)
    return tf.cast(x=n, dtype=kl.dtype) * kl
def _kl_sample(a, b, name='kl_sample'):
  """Batched KL divergence `KL(a || b)` for Sample distributions.

  We can leverage the fact that:

  ```
  KL(Sample(a) || Sample(b)) = sum(KL(a || b))
  ```

  where the sum is over the `sample_shape` dims.

  Args:
    a: Instance of `Sample` distribution.
    b: Instance of `Sample` distribution.
    name: (optional) name to use for created ops.
      Default value: `"kl_sample"`'.

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

  Raises:
    ValueError: If the `sample_shape` of `a` and `b` don't match.
  """
  assertions = []
  a_ss = tf.get_static_value(a.sample_shape)
  b_ss = tf.get_static_value(b.sample_shape)
  msg = '`a.sample_shape` must be identical to `b.sample_shape`.'
  if a_ss is not None and b_ss is not None:
    if not np.array_equal(a_ss, b_ss):
      raise ValueError(msg)
  elif a.validate_args or b.validate_args:
    assertions.append(assert_util.assert_equal(
        a.sample_shape, b.sample_shape, message=msg))
  with tf.control_dependencies(assertions):
    kl = kullback_leibler.kl_divergence(
        a.distribution, b.distribution, name=name)
    n = prefer_static.reduce_prod(a.sample_shape)
    return tf.cast(x=n, dtype=kl.dtype) * kl
def _kl_sample ( a , b , name = 'kl_sample' ) : assertions = [ ] a_ss = tf . get_static_value ( a . sample_shape ) b_ss = tf . get_static_value ( b . sample_shape ) msg = '`a.sample_shape` must be identical to `b.sample_shape`.' if a_ss is not None and b_ss is not None : if not np . array_equal ( a_ss , b_ss ) : raise ValueError ( msg ) elif a . validate_args or b . validate_args : assertions . append ( assert_util . assert_equal ( a . sample_shape , b . sample_shape , message = msg ) ) with tf . control_dependencies ( assertions ) : kl = kullback_leibler . kl_divergence ( a . distribution , b . distribution , name = name ) n = prefer_static . reduce_prod ( a . sample_shape ) return tf . cast ( x = n , dtype = kl . dtype ) * kl
Batched KL divergence `KL(a || b)` for Sample distributions.

  We can leverage the fact that:

  ```
  KL(Sample(a) || Sample(b)) = sum(KL(a || b))
  ```

  where the sum is over the `sample_shape` dims.

  Args:
    a: Instance of `Sample` distribution.
    b: Instance of `Sample` distribution.
    name: (optional) name to use for created ops.
      Default value: `"kl_sample"`'.

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

  Raises:
    ValueError: If the `sample_shape` of `a` and `b` don't match.
Batched KL divergence KL ( a || b ) for Sample distributions .