tensorflow_probability/python/layers/distribution_layer.py
def _make_kl_divergence_fn(
    distribution_b,
    use_exact_kl=False,
    test_points_reduce_axis=(),  # `None` == "all"; () == "none".
    test_points_fn=tf.convert_to_tensor,
    weight=None):
  """Creates a callable computing `KL[a,b]` from `a`, a `tfd.Distribution`."""

  if use_exact_kl is None:
    kl_divergence_fn = tfd.kl_divergence
  else:
    # Closure over: test_points_fn, test_points_reduce_axis.
    def kl_divergence_fn(distribution_a, distribution_b):
      z = test_points_fn(distribution_a)
      return tf.reduce_mean(
          input_tensor=distribution_a.log_prob(z) - distribution_b.log_prob(z),
          axis=test_points_reduce_axis)

  # Closure over: distribution_b, kl_divergence_fn, weight.
  def _fn(distribution_a):
    """Closure that computes KLDiv as a function of `a` as in `KL[a, b]`."""
    with tf.compat.v1.name_scope('kldivergence_loss'):
      # TODO(b/119756336): Due to eager/graph Jacobian graph caching bug
      # we add here the capability for deferred construction of the prior.
      # This capability can probably be removed once b/119756336 is resolved.
      distribution_b_ = (distribution_b() if callable(distribution_b)
                         else distribution_b)
      kl = kl_divergence_fn(distribution_a, distribution_b_)
      if weight is not None:
        kl = tf.cast(weight, dtype=kl.dtype) * kl
      # Losses appended with the model.add_loss and are expected to be a single
      # scalar, unlike model.loss, which is expected to be the loss per sample.
      # Therefore, we reduce over all dimensions, regardless of the shape.
      # We take the sum because (apparently) Keras will add this to the *post*
      # `reduce_sum` (total) loss.
      # TODO(b/126259176): Add end-to-end Keras/TFP test to ensure the API's
      # align, particularly wrt how losses are aggregated (across batch
      # members).
      return tf.reduce_sum(input_tensor=kl, name='batch_total_kl_divergence')

  return _fn
def _make_kl_divergence_fn(
    distribution_b,
    use_exact_kl=False,
    test_points_reduce_axis=(),  # `None` == "all"; () == "none".
    test_points_fn=tf.convert_to_tensor,
    weight=None):
  """Creates a callable computing `KL[a,b]` from `a`, a `tfd.Distribution`."""

  if use_exact_kl is None:
    kl_divergence_fn = tfd.kl_divergence
  else:
    # Closure over: test_points_fn, test_points_reduce_axis.
    def kl_divergence_fn(distribution_a, distribution_b):
      z = test_points_fn(distribution_a)
      return tf.reduce_mean(
          input_tensor=distribution_a.log_prob(z) - distribution_b.log_prob(z),
          axis=test_points_reduce_axis)

  # Closure over: distribution_b, kl_divergence_fn, weight.
  def _fn(distribution_a):
    """Closure that computes KLDiv as a function of `a` as in `KL[a, b]`."""
    with tf.compat.v1.name_scope('kldivergence_loss'):
      # TODO(b/119756336): Due to eager/graph Jacobian graph caching bug
      # we add here the capability for deferred construction of the prior.
      # This capability can probably be removed once b/119756336 is resolved.
      distribution_b_ = (distribution_b() if callable(distribution_b)
                         else distribution_b)
      kl = kl_divergence_fn(distribution_a, distribution_b_)
      if weight is not None:
        kl = tf.cast(weight, dtype=kl.dtype) * kl
      # Losses appended with the model.add_loss and are expected to be a single
      # scalar, unlike model.loss, which is expected to be the loss per sample.
      # Therefore, we reduce over all dimensions, regardless of the shape.
      # We take the sum because (apparently) Keras will add this to the *post*
      # `reduce_sum` (total) loss.
      # TODO(b/126259176): Add end-to-end Keras/TFP test to ensure the API's
      # align, particularly wrt how losses are aggregated (across batch
      # members).
      return tf.reduce_sum(input_tensor=kl, name='batch_total_kl_divergence')

  return _fn
def _make_kl_divergence_fn ( distribution_b , use_exact_kl = False , test_points_reduce_axis = ( ) , # `None` == "all"; () == "none". test_points_fn = tf . convert_to_tensor , weight = None ) : if use_exact_kl is None : kl_divergence_fn = tfd . kl_divergence else : # Closure over: test_points_fn, test_points_reduce_axis. def kl_divergence_fn ( distribution_a , distribution_b ) : z = test_points_fn ( distribution_a ) return tf . reduce_mean ( input_tensor = distribution_a . log_prob ( z ) - distribution_b . log_prob ( z ) , axis = test_points_reduce_axis ) # Closure over: distribution_b, kl_divergence_fn, weight. def _fn ( distribution_a ) : """Closure that computes KLDiv as a function of `a` as in `KL[a, b]`.""" with tf . compat . v1 . name_scope ( 'kldivergence_loss' ) : # TODO(b/119756336): Due to eager/graph Jacobian graph caching bug # we add here the capability for deferred construction of the prior. # This capability can probably be removed once b/119756336 is resolved. distribution_b_ = ( distribution_b ( ) if callable ( distribution_b ) else distribution_b ) kl = kl_divergence_fn ( distribution_a , distribution_b_ ) if weight is not None : kl = tf . cast ( weight , dtype = kl . dtype ) * kl # Losses appended with the model.add_loss and are expected to be a single # scalar, unlike model.loss, which is expected to be the loss per sample. # Therefore, we reduce over all dimensions, regardless of the shape. # We take the sum because (apparently) Keras will add this to the *post* # `reduce_sum` (total) loss. # TODO(b/126259176): Add end-to-end Keras/TFP test to ensure the API's # align, particularly wrt how losses are aggregated (across batch # members). return tf . reduce_sum ( input_tensor = kl , name = 'batch_total_kl_divergence' ) return _fn
Creates a callable computing `KL[a,b]` from `a`, a `tfd.Distribution`.
Creates a callable computing KL [ a b ] from a a tfd . Distribution .