code: def kl_divergence(
      distribution_a, distribution_b, allow_nan_stats=True, name=None
  ):
      """Get the KL-divergence KL(distribution_a || distribution_b).
  
      If there is no KL method registered specifically for `type(distribution_a)`
      and `type(distribution_b)`, then the class hierarchies of these types are
      searched.
  
      If one KL method is registered between any pairs of classes in these two
      parent hierarchies, it is used.
  
      If more than one such registered method exists, the method whose registered
      classes have the shortest sum MRO paths to the input types is used.
  
      If more than one such shortest path exists, the first method
      identified in the search is used (favoring a shorter MRO distance to
      `type(distribution_a)`).
  
      Args:
        distribution_a: The first distribution.
        distribution_b: The second distribution.
        allow_nan_stats: Python `bool`, default `True`. When `True`,
          statistics (e.g., mean, mode, variance) use the value "`NaN`" to
          indicate the result is undefined. When `False`, an exception is raised
          if one or more of the statistic's batch members are undefined.
        name: Python `str` name prefixed to Ops created by this class.
  
      Returns:
        A Tensor with the batchwise KL-divergence between `distribution_a`
        and `distribution_b`.
  
      Raises:
        NotImplementedError: If no KL method is defined for distribution types
          of `distribution_a` and `distribution_b`.
      """
      # NOTE: We use `d.__class__` instead of `type(d)` for objects that override
      # their `__class__` attribute.
      kl_fn = _registered_kl(
          distribution_a.__class__, distribution_b.__class__
      )
      if kl_fn is None:
          raise NotImplementedError(
              "No KL(distribution_a || distribution_b) registered for distribution_a "
              "type {} and distribution_b type {}".format(
                  distribution_a.__class__.__name__,
                  distribution_b.__class__.__name__,
              )
          )
  
      name = name or "KullbackLeibler"
      with tf.name_scope(name):
          # pylint: disable=protected-access
          with distribution_a._name_and_control_scope(name + "_a"):
              with distribution_b._name_and_control_scope(name + "_b"):
                  kl_t = kl_fn(
                      distribution_a, distribution_b, name=name
                  )
                  if allow_nan_stats:
                      return kl_t
  
          # Check KL for NaNs
          kl_t = tf.identity(kl_t, name="kl")
  
          with tf.control_dependencies(
              [
                  tf.debugging.Assert(
                      tf.logical_not(
                          tf.reduce_any(tf.math.is_nan(kl_t))
                      ),
                      [
                          (
                              "KL calculation between {} and {} returned NaN values "
                              "(and was called with allow_nan_stats=False). Values:".format(
                                  distribution_a.name,
                                  distribution_b.name,
                              )
                          ),
                          kl_t,
                      ],
                  )
              ]
          ):
              return tf.identity(kl_t, name="checked_kl")
  
file_path: probability/tensorflow_probability/python/distributions/kullback_leibler.py
imports: import inspect
  import six
  import tensorflow.compat.v2
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 2 for calling a global function with example usage: kl_divergence(distribution_a, distribution_b, allow_nan_stats, name) and returns: tf, kl_t