code: def _kl_masked_masked(a, b, name=None):
      """KL divergence between Masked distributions."""
      with tf.name_scope(name or "kl_masked_masked"):
          a_valid = tf.convert_to_tensor(a.validity_mask)
          b_valid = tf.convert_to_tensor(b.validity_mask)
          underlying_kl = kullback_leibler.kl_divergence(
              a.distribution, b.distribution
          )
  
          # The treatment for KL is as follows:
          # When both random variables are valid, the underlying KL applies.
          # When neither random variable is valid, the KL is 0., i.e.
          # `a log a - a log b = 0` because log a and log b are everywhere 0.
          # When exactly one is valid, we (a) raise an assertion error, if either
          # distribution's allow_nan_stats is set to False, or (b) return nan in
          # such positions.
          asserts = []
          if not (a.allow_nan_stats and b.allow_nan_stats):
              asserts.append(
                  assert_util.assert_equal(
                      a_valid,
                      b_valid,
                      message="KL is only valid for matching mask values",
                  )
              )
          with tf.control_dependencies(asserts):
              both_valid = a_valid & b_valid
              neither_valid = (~a_valid) & (~b_valid)
              dtype = underlying_kl.dtype
              return tf.where(
                  both_valid,
                  underlying_kl,
                  tf.where(
                      neither_valid, tf.zeros([], dtype), float("nan")
                  ),
              )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import batch_broadcast
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
label_desc: Use image node_id 4 for calling a global function with example usage: _kl_masked_masked(a, b, name) and returns: tf