//  [python]
def sample_ssl_sequence(sequence: tf.Tensor,
                        num_steps: int,
                        random: bool,
                        stride: int = 1,
                        num_windows: Optional[int] = 2) -> tf.Tensor:
  """Samples two segments of size num_steps randomly from a given sequence.

  Currently it only supports images, and specically designed for video self-
  supervised learning.

  Args:
    sequence: Any tensor where the first dimension is timesteps.
    num_steps: Number of steps (e.g. frames) to take.
    random: A boolean indicating whether to random sample the single window. If
      True, the offset is randomized. Only True is supported.
    stride: Distance to sample between timesteps.
    num_windows: Number of sequence sampled.

  Returns:
    A single Tensor with first dimension num_steps with the sampled segment.
  """
  sequence_length = tf.shape(sequence)[0]
  sequence_length = tf.cast(sequence_length, tf.float32)
  if random:
    max_offset = tf.cond(
        tf.greater(sequence_length, (num_steps - 1) * stride),
        lambda: sequence_length - (num_steps - 1) * stride,
        lambda: sequence_length)

    max_offset = tf.cast(max_offset, dtype=tf.float32)
    def cdf(k, power=1.0):
      """Cumulative distribution function for x^power."""
      p = -tf.math.pow(k, power + 1) / (
          power * tf.math.pow(max_offset, power + 1)) + k * (power + 1) / (
              power * max_offset)
      return p

    u = tf.random.uniform(())
    k_low = tf.constant(0, dtype=tf.float32)
    k_up = max_offset
    k = tf.math.floordiv(max_offset, 2.0)

    c = lambda k_low, k_up, k: tf.greater(tf.math.abs(k_up - k_low), 1.0)
    # pylint:disable=g-long-lambda
    b = lambda k_low, k_up, k: tf.cond(
        tf.greater(cdf(k), u),
        lambda: [k_low, k, tf.math.floordiv(k + k_low, 2.0)],
        lambda: [k, k_up, tf.math.floordiv(k_up + k, 2.0)])

    _, _, k = tf.while_loop(c, b, [k_low, k_up, k])
    delta = tf.cast(k, tf.int32)
    max_offset = tf.cast(max_offset, tf.int32)
    sequence_length = tf.cast(sequence_length, tf.int32)

    choice_1 = tf.cond(
        tf.equal(max_offset, sequence_length),
        lambda: tf.random.uniform((),
                                  maxval=tf.cast(max_offset, dtype=tf.int32),
                                  dtype=tf.int32),
        lambda: tf.random.uniform((),
                                  maxval=tf.cast(max_offset - delta,
                                                 dtype=tf.int32),
                                  dtype=tf.int32))
    choice_2 = tf.cond(
        tf.equal(max_offset, sequence_length),
        lambda: tf.random.uniform((),
                                  maxval=tf.cast(max_offset, dtype=tf.int32),
                                  dtype=tf.int32),
        lambda: choice_1 + delta)
    # pylint:disable=g-long-lambda
    shuffle_choice = tf.random.shuffle((choice_1, choice_2))
    offset_1 = shuffle_choice[0]
    offset_2 = shuffle_choice[1]

  else:
    raise NotImplementedError

  indices_1 = _sample_or_pad_sequence_indices(
      sequence=sequence,
      num_steps=num_steps,
      stride=stride,
      offset=offset_1)

  indices_2 = _sample_or_pad_sequence_indices(
      sequence=sequence,
      num_steps=num_steps,
      stride=stride,
      offset=offset_2)

  indices = tf.concat([indices_1, indices_2], axis=0)
  indices.set_shape((num_windows * num_steps,))
  output = tf.gather(sequence, indices)

  return output