code: def kolmogorov_smirnov_distance(
      samples, cdf, left_continuous_cdf=None, name=None
  ):
      """Computes the Kolmogorov-Smirnov distance between the given CDFs.
  
      The (absolute) Kolmogorov-Smirnov distance is the maximum (absolute)
      discrepancy between the CDFs, i.e.,
  
        sup_x(|cdf1(x) - cdf2(x)|)
  
      This is tractable to compute exactly when at least one CDF in question is an
      empirical CDF given by samples, because the analytic one need only be queried
      at the sampled values.
  
      If the distribution in question has atoms (e.g., is discrete), computing the
      distance requires CDF values for both sides of the discontinuity.  In this
      case, the `cdf` argument is assumed to compute the CDF inclusive of the atom,
      i.e., cdf(x) = Pr(X <= x).  The user must also supply the
      `left_continuous_cdf`, which must compute the cdf exclusive of the atom, i.e.,
      left_continuous_cdf(x) = Pr(X < x).
  
      For example, the two required cdfs of the degenerate distribution that places
      all the mass at 0 can be given as
      ```
      cdf=lambda x: tf.where(x < 0, 0., 1.)
      left_continuous_cdf=lambda x: tf.where(x <= 0, 0., 1.)
      ```
  
      Args:
        samples: Tensor of shape [n] + B.  Samples from some (batch of) scalar-event
          distribution(s) of interest, giving a (batch of) empirical CDF(s).
          Assumed IID across the 0 dimension.
        cdf: Analytic cdf inclusive of any atoms, as a function that can compute CDF
          values in batch.  Must accept a Tensor of shape [n] + B and the same dtype
          as `samples` and return a Tensor of shape [n] + B of CDF values.  For each
          sample x, `cdf(x) = Pr(X <= x)`.
        left_continuous_cdf: Analytic left-continuous cdf, as a function that can
          compute CDF values in batch.  Must accept a Tensor of shape [n] + B and
          the same dtype as `samples` and return a Tensor of shape [n] + B of CDF
          values.  For each sample x, `left_continuous_cdf(x) = Pr(X < x)`.  If the
          distribution under test has no atoms (i.e., the CDF is continuous), this
          is redundant and may be omitted.  Conversely, if this argument is omitted,
          the test assumes the distribution is atom-free.
        name: A name for this operation (optional).
  
      Returns:
        distance: Tensor of shape B: (Absolute) Kolmogorov-Smirnov distance between
          the empirical and analytic CDFs.
      """
      with tf.name_scope(name or "kolmogorov_smirnov_distance"):
          dtype = dtype_util.common_dtype([samples], tf.float32)
          samples = tf.convert_to_tensor(
              value=samples, name="samples", dtype=dtype
          )
          samples = tf.sort(samples, axis=0, direction="ASCENDING")
  
          # Compute analytic cdf values at each sample
          cdfs = cdf(samples)
          # Move the iid dimension of `cdfs` to the bottom so the empirical cdfs will
          # broadcast correctly
          cdfs = distribution_util.move_dimension(cdfs, 0, -1)
          if left_continuous_cdf is None:
              left_continuous_cdfs = cdfs
          else:
              left_continuous_cdfs = left_continuous_cdf(samples)
              left_continuous_cdfs = distribution_util.move_dimension(
                  left_continuous_cdfs, 0, -1
              )
  
          # Compute per-batch-member empirical cdf values at each sample
          # If any samples within a batch member are repeated, some of the entries
          # will be wrong:
          # - In low_empirical_cdfs, the first sample in a run of equal samples will
          #   have the correct cdf value, and the others will be too high; and
          # - In high_empirical_cdfs, the last sample in a run of equal samples will
          #   have the correct cdf value, and the others will be too low.
          # However, this is OK, because those errors do not change the maximums.
          # Could defensively use `empirical_cdfs` here, but those rely on the
          # relatively more expensive `searchsorted` operation.
          n = tf.cast(tf.shape(samples)[0], dtype=cdfs.dtype)
          low_empirical_cdfs = tf.range(n, dtype=cdfs.dtype) / n
          high_empirical_cdfs = tf.range(1, n + 1, dtype=cdfs.dtype) / n
  
          # Compute per-batch K-S distances on either side of each discontinuity in
          # the empirical CDF.  I only need one-sided comparisons in both cases,
          # because the empirical CDF is piecewise constant and the true CDF is
          # monotonic: The maximum of F(x) - F_n(x) occurs just before a
          # discontinuity, and the maximum of F_n(x) - F(x) occurs just after.
          low_distances = tf.reduce_max(
              left_continuous_cdfs - low_empirical_cdfs, axis=-1
          )
          high_distances = tf.reduce_max(
              high_empirical_cdfs - cdfs, axis=-1
          )
          return tf.maximum(low_distances, high_distances)
  
imports: import functools
  import itertools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.util.seed_stream import SeedStream
label_desc: Use image node_id 4 for calling a global function with example usage: kolmogorov_smirnov_distance(samples, cdf, left_continuous_cdf, name) and returns: tf