code: 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 = ps.reduce_prod(a.sample_shape)
          return tf.cast(x=n, dtype=kl.dtype) * kl
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  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 dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 2 for calling a global function with example usage: _kl_sample(a, b, name) and returns: unknown