code: def _make_kl_divergence_penalty(
      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:
          kl_divergence_fn = kullback_leibler.kl_divergence
      else:
  
          def kl_divergence_fn(distribution_a, distribution_b):
              z = test_points_fn(distribution_a)
              return tf.reduce_mean(
                  distribution_a.log_prob(z)
                  - distribution_b.log_prob(z),
                  axis=test_points_reduce_axis,
              )
  
      # Closure over: kl_divergence_fn, weight.
      def _fn(distribution_a, distribution_b):
          """Closure that computes KLDiv as a function of `a` as in `KL[a, b]`."""
          with tf.name_scope("kldivergence_loss"):
              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(kl, name="batch_total_kl_divergence")
  
      return _fn
  
file_path: probability/tensorflow_probability/python/layers/dense_variational_v2.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 1 for calling a global function with example usage: _make_kl_divergence_penalty(use_exact_kl, test_points_reduce_axis, test_points_fn, weight) and returns: _fn, tf, tf