//  [python]
def build_selfsup_transformations(num_flow_levels=3,
                                  seq_len=2,
                                  crop_height=0,
                                  crop_width=0,
                                  max_shift_height=0,
                                  max_shift_width=0,
                                  resize=True):
  """Apply augmentations to a list of student images."""

  def transform(images, i_or_ij, is_flow, crop_height, crop_width,
                shift_heights, shift_widths, resize):
    # Expect (i, j) for flows and masks and i for images.
    if isinstance(i_or_ij, int):
      i = i_or_ij
      # Flow needs i and j.
      assert not is_flow
    else:
      i, j = i_or_ij

    if is_flow:
      shifts = tf.stack([shift_heights, shift_widths], axis=-1)
      flow_offset = shifts[i] - shifts[j]
      images = images + tf.cast(flow_offset, tf.float32)

    shift_height = shift_heights[i]
    shift_width = shift_widths[i]
    height = images.shape[-3]
    width = images.shape[-2]

    # Assert that the cropped bounding box does not go out of the image frame.
    op1 = tf.compat.v1.assert_greater_equal(crop_height + shift_height, 0)
    op2 = tf.compat.v1.assert_greater_equal(crop_width + shift_width, 0)
    op3 = tf.compat.v1.assert_less_equal(height - crop_height + shift_height,
                                         height)
    op4 = tf.compat.v1.assert_less_equal(width - crop_width + shift_width,
                                         width)
    op5 = tf.compat.v1.assert_greater(
        height,
        2 * crop_height,
        message='Image height is too small for cropping.')
    op6 = tf.compat.v1.assert_greater(
        width, 2 * crop_width, message='Image width is too small for cropping.')
    with tf.control_dependencies([op1, op2, op3, op4, op5, op6]):
      images = images[:, crop_height + shift_height:height - crop_height +
                      shift_height, crop_width + shift_width:width -
                      crop_width + shift_width, :]
    if resize:
      images = uflow_utils.resize(images, height, width, is_flow=is_flow)
      images.set_shape((images.shape[0], height, width, images.shape[3]))
    else:
      images.set_shape((images.shape[0], height - 2 * crop_height,
                        width - 2 * crop_width, images.shape[3]))
    return images

  max_divisor = 2**(num_flow_levels - 1)
  assert crop_height % max_divisor == 0
  assert crop_width % max_divisor == 0
  assert max_shift_height <= crop_height
  assert max_shift_width <= crop_width
  # Compute random shifts for different images in a sequence.
  if max_shift_height > 0 or max_shift_width > 0:
    max_rand = max_shift_height // max_divisor
    shift_height_at_highest_level = tf.random.uniform([seq_len],
                                                      minval=-max_rand,
                                                      maxval=max_rand + 1,
                                                      dtype=tf.int32)
    shift_heights = shift_height_at_highest_level * max_divisor

    max_rand = max_shift_height // max_divisor
    shift_width_at_highest_level = tf.random.uniform([seq_len],
                                                     minval=-max_rand,
                                                     maxval=max_rand + 1,
                                                     dtype=tf.int32)
    shift_widths = shift_width_at_highest_level * max_divisor
  transform_fns = []
  for level in range(num_flow_levels):

    if max_shift_height == 0 and max_shift_width == 0:
      shift_heights = [0, 0]
      shift_widths = [0, 0]
    else:
      shift_heights = shift_heights // (2**level)
      shift_widths = shift_widths // (2**level)

    fn = partial(
        transform,
        crop_height=crop_height // (2**level),
        crop_width=crop_width // (2**level),
        shift_heights=shift_heights,
        shift_widths=shift_widths,
        resize=resize)
    transform_fns.append(fn)
  assert len(transform_fns) == num_flow_levels
  return transform_fns