jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __call__(self):
    def call(strategy=None):
      del strategy  # unused
      if self._period is None:
        return tf.no_op()
      if self._period == 1:
        return self._body()
      period = tf.cast(self._period, self._counter.dtype)
      remainder = tf.math.mod(self._counter.assign_add(1), period)
      return tf.cond(
          pred=tf.equal(remainder, 0), true_fn=self._body, false_fn=tf.no_op
      )

    # TODO(b/129083817) add an explicit unit test to ensure correct behavior
    ctx = tf.distribute.get_replica_context()
    if ctx:
      return tf.distribute.get_replica_context().merge_call(call)
    else:
      return call()
//  [python]
def get_episode_mask(time_steps):
  """Create a mask that is 0.0 for all final steps, 1.0 elsewhere.

  Args:
    time_steps: A TimeStep namedtuple representing a batch of steps.

  Returns:
    A float32 Tensor with 0s where step_type == LAST and 1s otherwise.
  """
  episode_mask = t.cast(
      tf.not_equal(time_steps.step_type, ts.StepType.LAST), tf.float32
  )
  return episode_mask
//  [python]
def get_episode_mask(time_steps):
  """Create a mask that is 0.0 for all final steps, 1.0 elsewhere.

  Args:
    time_steps: A TimeStep namedtuple representing a batch of steps.

  Returns:
    A float32 Tensor with 0s where step_type == LAST and 1s otherwise.
  """
  episode_mask = tf.cast(
      tf.not_equal(time_steps.step_type, ts.StepType.LAST), tf.float32
  )
  return episode_mask
//  [python]
def discounted_future_sum_masked(values, gamma, num_steps, episode_lengths):
  """Discounted future sum of batch-major values.

  Args:
    values: A Tensor of shape [batch_size, total_steps] and dtype float32.
    gamma: A float discount value.
    num_steps: A positive integer number of future steps to sum.
    episode_lengths: A vector shape [batch_size] with num_steps per episode.

  Returns:
    A Tensor of shape [batch_size, total_steps], where each entry is the
      discounted sum as in discounted_future_sum, except with values after
      the end of episode_lengths masked to 0.

  Raises:
    ValueError: If values is not of rank 2, or if total_steps is not defined.
  """
  if values.shape.rank != 2:
    raise ValueError('Input must be a rank 2 tensor.  Got %d.' % values.shape)

  total_steps = tf.compat.dimension_value(values.shape[1])
  if total_steps is None:
    raise ValueError(
        'total_steps dimension in input '
        'values[batch_size, total_steps] must be fully defined.'
    )

  episode_mask = tf.to_float(
      tf.sequence_mask(episode_lengths, total_steps)
  )
  values *= episode_mask
  return discounted_future_sum(values, gamma, num_steps)
//  [python]
def discounted_future_sum_masked(values, gamma, num_steps, episode_lengths):
  """Discounted future sum of batch-major values.

  Args:
    values: A Tensor of shape [batch_size, total_steps] and dtype float32.
    gamma: A float discount value.
    num_steps: A positive integer number of future steps to sum.
    episode_lengths: A vector shape [batch_size] with num_steps per episode.

  Returns:
    A Tensor of shape [batch_size, total_steps], where each entry is the
      discounted sum as in discounted_future_sum, except with values after
      the end of episode_lengths masked to 0.

  Raises:
    ValueError: If values is not of rank 2, or if total_steps is not defined.
  """
  if values.shape.rank != 2:
    raise ValueError('Input must be a rank 2 tensor.  Got %d.' % values.shape)

  total_steps = tf.compat.dimension_value(values.shape[1])
  if total_steps is None:
    raise ValueError(
        'total_steps dimension in input '
        'values[batch_size, total_steps] must be fully defined.'
    )

  episode_mask = tf.cast(
      tf.sequence_mask(episode_lengths, total_steps), tf.float32
  )
  values *= episode_mask
  return discounted_future_sum(values, gamma, num_steps)
//  [python]
def _get_gaussian_kernel(sigma, filter_shape):
  """Computes 1D Gaussian kernel."""
  sigma = tf.convert_to_tensor(sigma)
  x = tf.range(-filter_shape // 2 + 1, filter_shape // 2 + 1)
  x = tf.cast(x**2, sigma.dtype)
  x = tf.nn.softmax(-x / (2.0 * (sigma**2)))
  return x
//  [python]
def get_contiguous_sub_episodes(next_time_steps_discount):
  """Computes mask on sub-episodes which includes only contiguous components.

  Args:
    next_time_steps_discount: Tensor of shape [batch_size, total_steps]
      corresponding to environment discounts on next time steps (i.e.
      next_time_steps.discount).

  Returns:
    A float Tensor of shape [batch_size, total_steps] specifying mask including
      only contiguous components. Each row will be of the form
      [1.0] * a + [0.0] * b, where a >= 1 and b >= 0, and in which the initial
      sequence of ones corresponds to a contiguous sub-episode.
  """
  episode_end = tf.equal(
      next_time_steps_discount,
      tf.constant(0, dtype=next_time_steps_discount.dtype),
  )
  mask = tf.math.cumprod(
      1.0 - t.cast(episode_end, tf.float32), axis=1, exclusive=True
  )
  return mask
//  [python]
def get_contiguous_sub_episodes(next_time_steps_discount):
  """Computes mask on sub-episodes which includes only contiguous components.

  Args:
    next_time_steps_discount: Tensor of shape [batch_size, total_steps]
      corresponding to environment discounts on next time steps (i.e.
      next_time_steps.discount).

  Returns:
    A float Tensor of shape [batch_size, total_steps] specifying mask including
      only contiguous components. Each row will be of the form
      [1.0] * a + [0.0] * b, where a >= 1 and b >= 0, and in which the initial
      sequence of ones corresponds to a contiguous sub-episode.
  """
  episode_end = tf.equal(
      next_time_steps_discount,
      tf.constant(0, dtype=next_time_steps_discount.dtype),
  )
  mask = tf.math.cumprod(
      1.0 - tf.cast(episode_end, tf.float32), axis=1, exclusive=True
  )
  return mask
//  [python]
def solarize_add(image: tf.Tensor,
                 addition: int = 0,
                 threshold: int = 128) -> tf.Tensor:
  """Additive solarize the input image(s)."""
  # For each pixel in the image less than threshold
  # we add 'addition' amount to it and then clip the
  # pixel value to be between 0 and 255. The value
  # of 'addition' is between -128 and 128.
  added_image = t.cast(image, tf.int64) + addition
  added_image = t.cast(tf.clip_by_value(added_image, 0, 255), tf.uint8)
  return tf.where(image < threshold, added_image, image)
//  [python]
def solarize_add(image: tf.Tensor,
                 addition: int = 0,
                 threshold: int = 128) -> tf.Tensor:
  """Additive solarize the input image(s)."""
  # For each pixel in the image less than threshold
  # we add 'addition' amount to it and then clip the
  # pixel value to be between 0 and 255. The value
  # of 'addition' is between -128 and 128.
  added_image = tf.cast(image, tf.int64) + addition
  added_image = tf.cast(tf.clip_by_value(added_image, 0, 255), tf.uint8)
  return tf.where(image < threshold, added_image, image)
//  [python]
def discounted_future_sum(values, gamma, num_steps):
  """Discounted future sum of batch-major values.

  Args:
    values: A Tensor of shape [batch_size, total_steps] and dtype float32.
    gamma: A float discount value.
    num_steps: A positive integer number of future steps to sum.

  Returns:
    A Tensor of shape [batch_size, total_steps], where each entry `(i, j)` is
      the result of summing the entries of values starting from
      `gamma^0 * values[i, j]` to
      `gamma^(num_steps - 1) * values[i, j + num_steps - 1]`,
      with zeros padded to values.

      For example, values=[5, 6, 7], gamma=0.9, will result in sequence:
      
      [(5 * 0.9^0 + 6 * 0.9^1 + 7 * 0.9^2), (6 * 0.9^0 + 7 * 0.9^1), 7 * 0.9^0]
      

  Raises:
    ValueError: If values is not of rank 2.
  """
  if values.get_shape().rank != 2:
    raise ValueError(
        'Input must be rank 2 tensor.  Got %d.' % values.get_shape().rank
    )

  (batch_size, total_steps) = values.get_shape().as_list()

  num_steps = tf.minimum(num_steps, total_steps)
  discount_filter = tf.reshape(
      gamma ** tf.to_float(tf.range(num_steps)), [-1, 1, 1]
  )
  padded_values = tf.concat([values, tf.zeros([batch_size, num_steps - 1])], 1)

  convolved_values = tf.squeeze(
      tf.nn.conv1d(
          input=tf.expand_dims(padded_values, -1),
          filters=discount_filter,
          stride=1,
          padding='VALID',
      ),
      -1,
  )

  return convolved_values
//  [python]
def discounted_future_sum(values, gamma, num_steps):
  """Discounted future sum of batch-major values.

  Args:
    values: A Tensor of shape [batch_size, total_steps] and dtype float32.
    gamma: A float discount value.
    num_steps: A positive integer number of future steps to sum.

  Returns:
    A Tensor of shape [batch_size, total_steps], where each entry `(i, j)` is
      the result of summing the entries of values starting from
      `gamma^0 * values[i, j]` to
      `gamma^(num_steps - 1) * values[i, j + num_steps - 1]`,
      with zeros padded to values.

      For example, values=[5, 6, 7], gamma=0.9, will result in sequence:
      ```python
      [(5 * 0.9^0 + 6 * 0.9^1 + 7 * 0.9^2), (6 * 0.9^0 + 7 * 0.9^1), 7 * 0.9^0]
      ```

  Raises:
    ValueError: If values is not of rank 2.
  """
  if values.get_shape().rank != 2:
    raise ValueError(
        'Input must be rank 2 tensor.  Got %d.' % values.get_shape().rank
    )

  (batch_size, total_steps) = values.get_shape().as_list()

  num_steps = tf.minimum(num_steps, total_steps)
  discount_filter = tf.reshape(
      gamma ** tf.cast(tf.range(num_steps), tf.float32), [-1, 1, 1]
  )
  padded_values = tf.concat([values, tf.zeros([batch_size, num_steps - 1])], 1)

  convolved_values = tf.squeeze(
      tf.nn.conv1d(
          input=tf.expand_dims(padded_values, -1),
          filters=discount_filter,
          stride=1,
          padding='VALID',
      ),
      -1,
  )

  return convolved_values
//  [python]
def normalized_inputs(inputs):

    num_channels = inputs.get_shape()[-1]

    if inputs.get_shape().ndims != 4:
        raise ValueError('Input must be of size [batch_size, height, width, C>0]')

    if len(_CHANNEL_MEANS) != num_channels:
        raise ValueError('len(means) must match the number of channels')

    # We have a 1-D tensor of means; convert to 3-D.
    means_per_channel = tf.reshape(_CHANNEL_MEANS, [1, 1, num_channels])
    means_per_channel = tf.to_float(means_per_channel)

    stds_per_channel = tf.reshape(_CHANNEL_STDS, [1, 1, num_channels])
    stds_per_channel = tf.to_float(stds_per_channel)

    inputs = tf.subtract(inputs, means_per_channel)
    return tf.divide(inputs, stds_per_channel)
//  [python]
def normalized_inputs(inputs):

    num_channels = inputs.get_shape()[-1]

    if inputs.get_shape().ndims != 4:
        raise ValueError('Input must be of size [batch_size, height, width, C>0]')

    if len(_CHANNEL_MEANS) != num_channels:
        raise ValueError('len(means) must match the number of channels')

    # We have a 1-D tensor of means; convert to 3-D.
    means_per_channel = tf.reshape(_CHANNEL_MEANS, [1, 1, num_channels])
    means_per_channel = tf.cast(means_per_channel, dtype=inputs.dtype)

    stds_per_channel = tf.reshape(_CHANNEL_STDS, [1, 1, num_channels])
    stds_per_channel = tf.cast(stds_per_channel, dtype=inputs.dtype)

    inputs = tf.subtract(inputs, means_per_channel)
    return tf.divide(inputs, stds_per_channel)
//  [python]
def shift_values(values, gamma, num_steps, final_values=None):
  """Shifts batch-major values in time by some amount.

  Args:
    values: A Tensor of shape [batch_size, total_steps] and dtype float32.
    gamma: A float discount value.
    num_steps: A nonnegative integer amount to shift values by.
    final_values: A float32 Tensor of shape [batch_size] corresponding to the
      values at step num_steps + 1.  Defaults to None (all zeros).

  Returns:
    A Tensor of shape [batch_size, total_steps], where each entry (i, j) is
    gamma^num_steps * values[i, j + num_steps] if j + num_steps < total_steps;
    gamma^(total_steps - j) * final_values[i] otherwise.

  Raises:
    ValueError: If values is not of rank 2.
  """
  if values.get_shape().rank != 2:
    raise ValueError(
        'Input must be rank 2 tensor.  Got %d.' % values.get_shape().rank
    )

  (batch_size, total_steps) = values.get_shape().as_list()
  num_steps = tf.minimum(num_steps, total_steps)

  if final_values is None:
    final_values = tf.zeros([batch_size])

  padding_exponent = tf.expand_dims(
      tf.to_float(tf.range(num_steps, 0, -1)), 0
  )
  final_pad = tf.expand_dims(final_values, 1) * gamma**padding_exponent
  return tf.concat(
      [
          gamma ** tf.to_float(num_steps) * values[:, num_steps:],
          final_pad,
      ],
      1,
  )
//  [python]
def shift_values(values, gamma, num_steps, final_values=None):
  """Shifts batch-major values in time by some amount.

  Args:
    values: A Tensor of shape [batch_size, total_steps] and dtype float32.
    gamma: A float discount value.
    num_steps: A nonnegative integer amount to shift values by.
    final_values: A float32 Tensor of shape [batch_size] corresponding to the
      values at step num_steps + 1.  Defaults to None (all zeros).

  Returns:
    A Tensor of shape [batch_size, total_steps], where each entry (i, j) is
    gamma^num_steps * values[i, j + num_steps] if j + num_steps < total_steps;
    gamma^(total_steps - j) * final_values[i] otherwise.

  Raises:
    ValueError: If values is not of rank 2.
  """
  if values.get_shape().rank != 2:
    raise ValueError(
        'Input must be rank 2 tensor.  Got %d.' % values.get_shape().rank
    )

  (batch_size, total_steps) = values.get_shape().as_list()
  num_steps = tf.minimum(num_steps, total_steps)

  if final_values is None:
    final_values = tf.zeros([batch_size])

  padding_exponent = tf.expand_dims(
      tf.cast(tf.range(num_steps, 0, -1), tf.float32), 0
  )
  final_pad = tf.expand_dims(final_values, 1) * gamma**padding_exponent
  return tf.concat(
      [
          gamma ** tf.cast(num_steps, tf.float32) * values[:, num_steps:],
          final_pad,
      ],
      1,
  )
//  [python]
def distort(self, image: tf.Tensor) -> tf.Tensor:
    """See base class."""
    input_image_type = image.dtype
    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = t.cast(image, dtype=tf.uint8)

    tf_policies = self._make_tf_policies()
    image, _ = select_and_apply_random_policy(tf_policies, image, bboxes=None)
    image = t.cast(image, dtype=input_image_type)
    return image
//  [python]
def distort(self, image: tf.Tensor) -> tf.Tensor:
    """See base class."""
    input_image_type = image.dtype
    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = tf.cast(image, dtype=tf.uint8)

    tf_policies = self._make_tf_policies()
    image, _ = select_and_apply_random_policy(tf_policies, image, bboxes=None)
    image = tf.cast(image, dtype=input_image_type)
    return image
//  [python]
def _apply_func_with_prob(func: Any, image: tf.Tensor,
                          bboxes: Optional[tf.Tensor], args: Any, prob: float):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)
  assert inspect.getfullargspec(func)[0][1] == 'bboxes'

  # Apply the function with probability `prob`.
  should_apply_op = tf.floor(tf.random.uniform([], dtype=tf.float32) + prob) > 0
  augmented_image, augmented_bboxes = tf.cond(
      should_apply_op,
      lambda: func(image, bboxes, *args),
      lambda: (image, bboxes))
  return augmented_image, augmented_bboxes
//  [python]
def _apply_func_with_prob(func: Any, image: tf.Tensor,
                          bboxes: Optional[tf.Tensor], args: Any, prob: float):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)
  assert inspect.getfullargspec(func)[0][1] == 'bboxes'

  # Apply the function with probability `prob`.
  should_apply_op = tf.cast(
      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)
  augmented_image, augmented_bboxes = tf.cond(
      should_apply_op,
      lambda: func(image, bboxes, *args),
      lambda: (image, bboxes))
  return augmented_image, augmented_bboxes
//  [python]
def distort_with_boxes(self, image: tf.Tensor,
                         bboxes: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    """See base class."""
    input_image_type = image.dtype
    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = tf.cast(image, dtype=tf.uint8)

    tf_policies = self._make_tf_policies()
    image, bboxes = select_and_apply_random_policy(tf_policies, image, bboxes)
    image = tf.cast(image, dtype=input_image_type)
    assert bboxes is not None
    return image, bboxes
//  [python]
def _apply_bbox_augmentation_wrapper(image, bbox, new_bboxes, prob,
                                     augmentation_func, func_changes_bbox,
                                     *args):
  """Applies _apply_bbox_augmentation with probability prob.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    new_bboxes: 2D Tensor that is a list of the bboxes in the image after they
      have been altered by aug_func. These will only be changed when
      func_changes_bbox is set to true. Each bbox has 4 elements
      (min_y, min_x, max_y, max_x) of type float that are the normalized
      bbox coordinates between 0 and 1.
    prob: Float that is the probability of applying _apply_bbox_augmentation.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    func_changes_bbox: Boolean. Does augmentation_func return bbox in addition
      to image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A tuple. Fist element is a modified version of image, where the bbox
    location in the image will have augmentation_func applied to it if it is
    chosen to be called with probability `prob`. The second element is a
    Tensor of Tensors of length 4 that will contain the altered bbox after
    applying augmentation_func.
  """
  should_apply_op = tf.floor(tf.random.uniform([], dtype=tf.float32) + prob)
  should_apply_op = tf.cast(should_apply_op, tf.bool)
  if func_changes_bbox:
    augmented_image, bbox = tf.cond(
        should_apply_op,
        lambda: augmentation_func(image, bbox, *args),
        lambda: (image, bbox))
  else:
    augmented_image = tf.cond(
        should_apply_op,
        lambda: _apply_bbox_augmentation(image, bbox, augmentation_func, *args),
        lambda: image)
  new_bboxes = _concat_bbox(bbox, new_bboxes)
  return augmented_image, new_bboxes
//  [python]
def _apply_bbox_augmentation_wrapper(image, bbox, new_bboxes, prob,
                                     augmentation_func, func_changes_bbox,
                                     *args):
  """Applies _apply_bbox_augmentation with probability prob.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    new_bboxes: 2D Tensor that is a list of the bboxes in the image after they
      have been altered by aug_func. These will only be changed when
      func_changes_bbox is set to true. Each bbox has 4 elements
      (min_y, min_x, max_y, max_x) of type float that are the normalized
      bbox coordinates between 0 and 1.
    prob: Float that is the probability of applying _apply_bbox_augmentation.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    func_changes_bbox: Boolean. Does augmentation_func return bbox in addition
      to image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A tuple. Fist element is a modified version of image, where the bbox
    location in the image will have augmentation_func applied to it if it is
    chosen to be called with probability `prob`. The second element is a
    Tensor of Tensors of length 4 that will contain the altered bbox after
    applying augmentation_func.
  """
  should_apply_op = tf.cast(
      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)
  if func_changes_bbox:
    augmented_image, bbox = tf.cond(
        should_apply_op,
        lambda: augmentation_func(image, bbox, *args),
        lambda: (image, bbox))
  else:
    augmented_image = tf.cond(
        should_apply_op,
        lambda: _apply_bbox_augmentation(image, bbox, augmentation_func, *args),
        lambda: image)
  new_bboxes = _concat_bbox(bbox, new_bboxes)
  return augmented_image, new_bboxes
//  [python]
def _apply_bbox_augmentation(image, bbox, augmentation_func, *args):
  """Applies augmentation_func to the subsection of image indicated by bbox.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A modified version of image, where the bbox location in the image will
    have `ugmentation_func applied to it.
  """
  image_height = tf.cast(tf.shape(image)[0], tf.float32)
  image_width = tf.cast(tf.shape(image)[1], tf.float32)
  min_y = tf.cast(image_height * bbox[0], tf.int32)
  min_x = tf.cast(image_width * bbox[1], tf.int32)
  max_y = tf.cast(image_height * bbox[2], tf.int32)
  max_x = tf.cast(image_width * bbox[3], tf.int32)
  image_height = tf.cast(image_height, tf.int32)
  image_width = tf.cast(image_width, tf.int32)

  # Clip to be sure the max values do not fall out of range.
  max_y = tf.minimum(max_y, image_height - 1)
  max_x = tf.minimum(max_x, image_width - 1)

  # Get the sub-tensor that is the image within the bounding box region.
  bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

  # Apply the augmentation function to the bbox portion of the image.
  augmented_bbox_content = augmentation_func(bbox_content, *args)

  # Pad the augmented_bbox_content and the mask to match the shape of original
  # image.
  augmented_bbox_content = tf.pad(augmented_bbox_content,
                                  [[min_y, (image_height - 1) - max_y],
                                   [min_x, (image_width - 1) - max_x],
                                   [0, 0]])

  # Create a mask that will be used to zero out a part of the original image.
  mask_tensor = tf.zeros_like(bbox_content)

  mask_tensor = tf.pad(mask_tensor,
                       [[min_y, (image_height - 1) - max_y],
                        [min_x, (image_width - 1) - max_x],
                        [0, 0]],
                       constant_values=1)
  # Replace the old bbox content with the new augmented content.
  image = image * mask_tensor + augmented_bbox_content
  return image
//  [python]
def _shift_bbox(bbox, image_height, image_width, pixels, shift_horizontal):
  """Shifts the bbox coordinates by pixels.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, width of the image.
    pixels: An int. How many pixels to shift the bbox.
    shift_horizontal: Boolean. If true then shift in X dimension else shift in
      Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  pixels = t.cast(pixels, t.int32)
  # Convert bbox to integer pixel locations.
  min_y = t.cast(t.cast(image_height, t.float32) * bbox[0], t.int32)
  min_x = t.cast(t.cast(image_width, t.float32) * bbox[1], t.int32)
  max_y = t.cast(t.cast(image_height, t.float32) * bbox[2], t.int32)
  max_x = t.cast(t.cast(image_width, t.float32) * bbox[3], t.int32)

  if shift_horizontal:
    min_x = t.maximum(0, min_x - pixels)
    max_x = t.minimum(image_width, max_x - pixels)
  else:
    min_y = t.maximum(0, min_y - pixels)
    max_y = t.minimum(image_height, max_y - pixels)

  # Convert bbox back to floats.
  min_y = t.cast(min_y, t.float32) / t.cast(image_height, t.float32)
  min_x = t.cast(min_x, t.float32) / t.cast(image_width, t.float32)
  max_y = t.cast(max_y, t.float32) / t.cast(image_height, t.float32)
  max_x = t.cast(max_x, t.float32) / t.cast(image_width, t.float32)

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return t.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shift_bbox(bbox, image_height, image_width, pixels, shift_horizontal):
  """Shifts the bbox coordinates by pixels.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, width of the image.
    pixels: An int. How many pixels to shift the bbox.
    shift_horizontal: Boolean. If true then shift in X dimension else shift in
      Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  pixels = tf.cast(pixels, tf.int32)
  # Convert bbox to integer pixel locations.
  min_y = tf.cast(tf.cast(image_height, tf.float32) * bbox[0], tf.int32)
  min_x = tf.cast(tf.cast(image_width, tf.float32) * bbox[1], tf.int32)
  max_y = tf.cast(tf.cast(image_height, tf.float32) * bbox[2], tf.int32)
  max_x = tf.cast(tf.cast(image_width, tf.float32) * bbox[3], tf.int32)

  if shift_horizontal:
    min_x = tf.maximum(0, min_x - pixels)
    max_x = tf.minimum(image_width, max_x - pixels)
  else:
    min_y = tf.maximum(0, min_y - pixels)
    max_y = tf.minimum(image_height, max_y - pixels)

  # Convert bbox back to floats.
  min_y = tf.cast(min_y, tf.float32) / tf.cast(image_height, tf.float32)
  min_x = tf.cast(min_x, tf.float32) / tf.cast(image_width, tf.float32)
  max_y = tf.cast(max_y, tf.float32) / tf.cast(image_height, tf.float32)
  max_x = tf.cast(max_x, tf.float32) / tf.cast(image_width, tf.float32)

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _mixup(self, images: tf.Tensor,
             labels: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
    """Applies mixup."""
    lam = MixupAndCutmix._sample_from_beta(self.mixup_alpha, self.mixup_alpha,
                                           tf.shape(labels))
    if images.shape.rank == 4:
      lam = tf.reshape(lam, [-1, 1, 1, 1])
    elif images.shape.rank == 5:
      lam = tf.reshape(lam, [-1, 1, 1, 1, 1])
    else:
      raise ValueError('Bad image rank: {}'.format(images.shape.rank))

    lam_cast = tf.to_float(lam)
    images = lam_cast * images + (1. - lam_cast) * tf.reverse(images, [0])

    return images, labels, tf.squeeze(lam)
//  [python]
def _mixup(self, images: tf.Tensor,
             labels: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
    """Applies mixup."""
    lam = MixupAndCutmix._sample_from_beta(self.mixup_alpha, self.mixup_alpha,
                                           tf.shape(labels))
    if images.shape.rank == 4:
      lam = tf.reshape(lam, [-1, 1, 1, 1])
    elif images.shape.rank == 5:
      lam = tf.reshape(lam, [-1, 1, 1, 1, 1])
    else:
      raise ValueError('Bad image rank: {}'.format(images.shape.rank))

    lam_cast = tf.cast(lam, dtype=images.dtype)
    images = lam_cast * images + (1. - lam_cast) * tf.reverse(images, [0])

    return images, labels, tf.squeeze(lam)
//  [python]
def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
  """Shifts the bbox according to how the image was sheared.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    level: Float. How much to shear the image.
    shear_horizontal: If true then shear in X dimension else shear in
      the Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  image_height, image_width = (
      t.cast(image_height, t.float32), t.cast(image_width, t.float32))

  # Change bbox coordinates to be pixels.
  min_y = t.cast(image_height * bbox[0], t.int32)
  min_x = t.cast(image_width * bbox[1], t.int32)
  max_y = t.cast(image_height * bbox[2], t.int32)
  max_x = t.cast(image_width * bbox[3], t.int32)
  coordinates = t.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = t.cast(coordinates, t.float32)

  # Shear the coordinates according to the translation matrix.
  if shear_horizontal:
    translation_matrix = t.stack(
        [[1, 0], [-level, 1]])
  else:
    translation_matrix = t.stack(
        [[1, -level], [0, 1]])
  translation_matrix = t.cast(translation_matrix, t.float32)
  new_coords = t.cast(
      t.matmul(translation_matrix, t.transpose(coordinates)), t.int32)

  # Find min/max values and convert them back to floats.
  min_y = t.cast(t.reduce_min(new_coords[0, :]), t.float32) / image_height
  min_x = t.cast(t.reduce_min(new_coords[1, :]), t.float32) / image_width
  max_y = t.cast(t.reduce_max(new_coords[0, :]), t.float32) / image_height
  max_x = t.cast(t.reduce_max(new_coords[1, :]), t.float32) / image_width

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return t.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
  """Shifts the bbox according to how the image was sheared.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    level: Float. How much to shear the image.
    shear_horizontal: If true then shear in X dimension else shear in
      the Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  image_height, image_width = (
      tf.cast(image_height, tf.float32), tf.cast(image_width, tf.float32))

  # Change bbox coordinates to be pixels.
  min_y = tf.cast(image_height * bbox[0], tf.int32)
  min_x = tf.cast(image_width * bbox[1], tf.int32)
  max_y = tf.cast(image_height * bbox[2], tf.int32)
  max_x = tf.cast(image_width * bbox[3], tf.int32)
  coordinates = tf.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = tf.cast(coordinates, tf.float32)

  # Shear the coordinates according to the translation matrix.
  if shear_horizontal:
    translation_matrix = tf.stack(
        [[1, 0], [-level, 1]])
  else:
    translation_matrix = tf.stack(
        [[1, -level], [0, 1]])
  translation_matrix = tf.cast(translation_matrix, tf.float32)
  new_coords = tf.cast(
      tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

  # Find min/max values and convert them back to floats.
  min_y = tf.cast(tf.reduce_min(new_coords[0, :]), tf.float32) / image_height
  min_x = tf.cast(tf.reduce_min(new_coords[1, :]), tf.float32) / image_width
  max_y = tf.cast(tf.reduce_max(new_coords[0, :]), tf.float32) / image_height
  max_x = tf.cast(tf.reduce_max(new_coords[1, :]), tf.float32) / image_width

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _rotate_bbox(bbox, image_height, image_width, degrees):
  """Rotates the bbox coordinated by degrees.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    degrees: Float, a scalar angle in degrees to rotate all images by. If
      degrees is positive the image will be rotated clockwise otherwise it will
      be rotated counterclockwise.

  Returns:
    A tensor of the same shape as bbox, but now with the rotated coordinates.
  """
  image_height, image_width = (
      t.cast(image_height, t.float32), t.cast(image_width, t.float32))

  # Convert from degrees to radians.
  degrees_to_radians = math.pi / 180.0
  radians = degrees * degrees_to_radians

  # Translate the bbox to the center of the image and turn the normalized 0-1
  # coordinates to absolute pixel locations.
  # Y coordinates are made negative as the y axis of images goes down with
  # increasing pixel values, so we negate to make sure x axis and y axis points
  # are in the traditionally positive direction.
  min_y = -t.cast(image_height * (bbox[0] - 0.5), t.int32)
  min_x = t.cast(image_width * (bbox[1] - 0.5), t.int32)
  max_y = -t.cast(image_height * (bbox[2] - 0.5), t.int32)
  max_x = t.cast(image_width * (bbox[3] - 0.5), t.int32)
  coordinates = t.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = t.cast(coordinates, t.float32)
  # Rotate the coordinates according to the rotation matrix clockwise if
  # radians is positive, else negative
  rotation_matrix = t.stack(
      [[t.cos(radians), t.sin(radians)],
       [-t.sin(radians), t.cos(radians)]])
  new_coords = t.cast(
      t.matmul(rotation_matrix, t.transpose(coordinates)), t.int32)
  # Find min/max values and convert them back to normalized 0-1 floats.
  min_y = -(
      t.cast(t.reduce_max(new_coords[0, :]), t.float32) / image_height - 0.5)
  min_x = t.cast(t.reduce_min(new_coords[1, :]),
                  t.float32) / image_width + 0.5
  max_y = -(
      t.cast(t.reduce_min(new_coords[0, :]), t.float32) / image_height - 0.5)
  max_x = t.cast(t.reduce_max(new_coords[1, :]),
                  t.float32) / image_width + 0.5

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return t.stack([min_y, min_x, max_y, max_x])
//  [python]
def _rotate_bbox(bbox, image_height, image_width, degrees):
  """Rotates the bbox coordinated by degrees.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    degrees: Float, a scalar angle in degrees to rotate all images by. If
      degrees is positive the image will be rotated clockwise otherwise it will
      be rotated counterclockwise.

  Returns:
    A tensor of the same shape as bbox, but now with the rotated coordinates.
  """
  image_height, image_width = (
      tf.cast(image_height, tf.float32), tf.cast(image_width, tf.float32))

  # Convert from degrees to radians.
  degrees_to_radians = math.pi / 180.0
  radians = degrees * degrees_to_radians

  # Translate the bbox to the center of the image and turn the normalized 0-1
  # coordinates to absolute pixel locations.
  # Y coordinates are made negative as the y axis of images goes down with
  # increasing pixel values, so we negate to make sure x axis and y axis points
  # are in the traditionally positive direction.
  min_y = -tf.cast(image_height * (bbox[0] - 0.5), tf.int32)
  min_x = tf.cast(image_width * (bbox[1] - 0.5), tf.int32)
  max_y = -tf.cast(image_height * (bbox[2] - 0.5), tf.int32)
  max_x = tf.cast(image_width * (bbox[3] - 0.5), tf.int32)
  coordinates = tf.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = tf.cast(coordinates, tf.float32)
  # Rotate the coordinates according to the rotation matrix clockwise if
  # radians is positive, else negative
  rotation_matrix = tf.stack(
      [[tf.cos(radians), tf.sin(radians)],
       [-tf.sin(radians), tf.cos(radians)]])
  new_coords = tf.cast(
      tf.matmul(rotation_matrix, tf.transpose(coordinates)), tf.int32)
  # Find min/max values and convert them back to normalized 0-1 floats.
  min_y = -(
      tf.cast(tf.reduce_max(new_coords[0, :]), tf.float32) / image_height - 0.5)
  min_x = tf.cast(tf.reduce_min(new_coords[1, :]),
                  tf.float32) / image_width + 0.5
  max_y = -(
      tf.cast(tf.reduce_min(new_coords[0, :]), tf.float32) / image_height - 0.5)
  max_x = tf.cast(tf.reduce_max(new_coords[1, :]),
                  tf.float32) / image_width + 0.5

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def filter_boxes_by_ioa(
    bboxes: tf.Tensor, crop_box: tf.Tensor, min_box_overlap: float
) -> tf.Tensor:
  """Filter boxes by intersection over area (IOA).

  The boxes with IOA less than min_box_overlap will be replaced by
  (0, 0, 0, 0) so they can be filtered out later.

  Args:
    bboxes: a float tensor of shape [N, 4] representing normalized bounding box
      coordinates.
    crop_box: a float tensor of shape [1, 1, 4] representing the normalized crop
      box.
    min_box_overlap: minimum overlap of the box with the crop box to keep the
      box.

  Returns:
    a tensor of shape [N, 4] with filtered box coordinates replaced by 0.
  """
  ioas = box_ops.bbox_intersection_over_area(bboxes[None, ...], crop_box)[0]
  keep_boxes = ioas >= min_box_overlap
  # Set coordinates to (0, 0, 0, 0) for filtered boxes
  return bboxes * tf.to_float(keep_boxes)
//  [python]
def filter_boxes_by_ioa(
    bboxes: tf.Tensor, crop_box: tf.Tensor, min_box_overlap: float
) -> tf.Tensor:
  """Filter boxes by intersection over area (IOA).

  The boxes with IOA less than min_box_overlap will be replaced by
  (0, 0, 0, 0) so they can be filtered out later.

  Args:
    bboxes: a float tensor of shape [N, 4] representing normalized bounding box
      coordinates.
    crop_box: a float tensor of shape [1, 1, 4] representing the normalized crop
      box.
    min_box_overlap: minimum overlap of the box with the crop box to keep the
      box.

  Returns:
    a tensor of shape [N, 4] with filtered box coordinates replaced by 0.
  """
  ioas = box_ops.bbox_intersection_over_area(bboxes[None, ...], crop_box)[0]
  keep_boxes = ioas >= min_box_overlap
  # Set coordinates to (0, 0, 0, 0) for filtered boxes
  return bboxes * tf.cast(keep_boxes, dtype=bboxes.dtype)
//  [python]
def _cutmix(self, images: tf.Tensor,
              labels: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
    """Applies cutmix."""
    lam = MixupAndCutmix._sample_from_beta(self.cutmix_alpha, self.cutmix_alpha,
                                           tf.shape(labels))

    ratio = tf.math.sqrt(1 - lam)

    batch_size = tf.shape(images)[0]

    if images.shape.rank == 4:
      image_height, image_width = tf.shape(images)[1], tf.shape(images)[2]
      fill_fn = _fill_rectangle
    elif images.shape.rank == 5:
      image_height, image_width = tf.shape(images)[2], tf.shape(images)[3]
      fill_fn = _fill_rectangle_video
    else:
      raise ValueError('Bad image rank: {}'.format(images.shape.rank))

    cut_height = tf.to_int32(
        ratio * tf.to_float(image_height))
    cut_width = tf.to_int32(
        ratio * tf.to_float(image_height))

    random_center_height = tf.random.uniform(
        shape=[batch_size], minval=0, maxval=image_height, dtype=tf.int32)
    random_center_width = tf.random.uniform(
        shape=[batch_size], minval=0, maxval=image_width, dtype=tf.int32)

    bbox_area = cut_height * cut_width
    lam = 1. - bbox_area / (image_height * image_width)
    lam = tf.to_float(lam)

    images = tf.map_fn(
        lambda x: fill_fn(*x),
        (images, random_center_width, random_center_height, cut_width // 2,
         cut_height // 2, tf.reverse(images, [0])),
        dtype=(
            images.dtype, tf.int32, tf.int32, tf.int32, tf.int32, images.dtype),
        fn_output_signature=tf.TensorSpec(images.shape[1:], dtype=images.dtype))

    return images, labels, lam
//  [python]
def _cutmix(self, images: tf.Tensor,
              labels: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
    """Applies cutmix."""
    lam = MixupAndCutmix._sample_from_beta(self.cutmix_alpha, self.cutmix_alpha,
                                           tf.shape(labels))

    ratio = tf.math.sqrt(1 - lam)

    batch_size = tf.shape(images)[0]

    if images.shape.rank == 4:
      image_height, image_width = tf.shape(images)[1], tf.shape(images)[2]
      fill_fn = _fill_rectangle
    elif images.shape.rank == 5:
      image_height, image_width = tf.shape(images)[2], tf.shape(images)[3]
      fill_fn = _fill_rectangle_video
    else:
      raise ValueError('Bad image rank: {}'.format(images.shape.rank))

    cut_height = tf.cast(
        ratio * tf.cast(image_height, dtype=tf.float32), dtype=tf.int32)
    cut_width = tf.cast(
        ratio * tf.cast(image_height, dtype=tf.float32), dtype=tf.int32)

    random_center_height = tf.random.uniform(
        shape=[batch_size], minval=0, maxval=image_height, dtype=tf.int32)
    random_center_width = tf.random.uniform(
        shape=[batch_size], minval=0, maxval=image_width, dtype=tf.int32)

    bbox_area = cut_height * cut_width
    lam = 1. - bbox_area / (image_height * image_width)
    lam = tf.cast(lam, dtype=tf.float32)

    images = tf.map_fn(
        lambda x: fill_fn(*x),
        (images, random_center_width, random_center_height, cut_width // 2,
         cut_height // 2, tf.reverse(images, [0])),
        dtype=(
            images.dtype, tf.int32, tf.int32, tf.int32, tf.int32, images.dtype),
        fn_output_signature=tf.TensorSpec(images.shape[1:], dtype=images.dtype))

    return images, labels, lam
//  [python]
def LOG(x, m, *a):
  xf = t.cast(x, t.float32)
  print_op = tf.print(
      m, '(',
      tf.shape(xf),
      tf.reduce_min(xf),
      tf.reduce_mean(xf),
      tf.reduce_max(xf),
      ')', *a)
  with tf.control_dependencies([print_op]):
    return tf.identity(x)
//  [python]
def LOG(x, m, *a):
  xf = tf.cast(x, tf.float32)
  print_op = tf.print(
      m, '(',
      tf.shape(xf),
      tf.reduce_min(xf),
      tf.reduce_mean(xf),
      tf.reduce_max(xf),
      ')', *a)
  with tf.control_dependencies([print_op]):
    return tf.identity(x)
//  [python]
def _erase(self, image: tf.Tensor) -> tf.Tensor:
    """Erase an area."""
    if self._min_count == self._max_count:
      count = self._min_count
    else:
      count = tf.random.uniform(
          shape=[],
          minval=int(self._min_count),
          maxval=int(self._max_count - self._min_count + 1),
          dtype=tf.int32)

    image_height = tf.shape(image)[0]
    image_width = tf.shape(image)[1]
    area = tf.to_float(image_width * image_height)

    for _ in range(count):
      # Work around since break is not supported in tf.function
      is_trial_successfull = False
      for _ in range(self._trials):
        if not is_trial_successfull:
          erase_area = tf.random.uniform(
              shape=[],
              minval=area * self._min_area,
              maxval=area * self._max_area)
          aspect_ratio = tf.math.exp(
              tf.random.uniform(
                  shape=[],
                  minval=self._min_log_aspect,
                  maxval=self._max_log_aspect))

          half_height = tf.to_int32(
              tf.math.round(tf.math.sqrt(erase_area * aspect_ratio) / 2))
          half_width = tf.to_int32(
              tf.math.round(tf.math.sqrt(erase_area / aspect_ratio) / 2))

          if 2 * half_height < image_height and 2 * half_width < image_width:
            center_height = tf.random.uniform(
                shape=[],
                minval=0,
                maxval=int(image_height - 2 * half_height),
                dtype=tf.int32)
            center_width = tf.random.uniform(
                shape=[],
                minval=0,
                maxval=int(image_width - 2 * half_width),
                dtype=tf.int32)

            image = _fill_rectangle(
                image,
                center_width,
                center_height,
                half_width,
                half_height,
                replace=None)

            is_trial_successfull = True

    return image
//  [python]
def _erase(self, image: tf.Tensor) -> tf.Tensor:
    """Erase an area."""
    if self._min_count == self._max_count:
      count = self._min_count
    else:
      count = tf.random.uniform(
          shape=[],
          minval=int(self._min_count),
          maxval=int(self._max_count - self._min_count + 1),
          dtype=tf.int32)

    image_height = tf.shape(image)[0]
    image_width = tf.shape(image)[1]
    area = tf.cast(image_width * image_height, tf.float32)

    for _ in range(count):
      # Work around since break is not supported in tf.function
      is_trial_successfull = False
      for _ in range(self._trials):
        if not is_trial_successfull:
          erase_area = tf.random.uniform(
              shape=[],
              minval=area * self._min_area,
              maxval=area * self._max_area)
          aspect_ratio = tf.math.exp(
              tf.random.uniform(
                  shape=[],
                  minval=self._min_log_aspect,
                  maxval=self._max_log_aspect))

          half_height = tf.cast(
              tf.math.round(tf.math.sqrt(erase_area * aspect_ratio) / 2),
              dtype=tf.int32)
          half_width = tf.cast(
              tf.math.round(tf.math.sqrt(erase_area / aspect_ratio) / 2),
              dtype=tf.int32)

          if 2 * half_height < image_height and 2 * half_width < image_width:
            center_height = tf.random.uniform(
                shape=[],
                minval=0,
                maxval=int(image_height - 2 * half_height),
                dtype=tf.int32)
            center_width = tf.random.uniform(
                shape=[],
                minval=0,
                maxval=int(image_width - 2 * half_width),
                dtype=tf.int32)

            image = _fill_rectangle(
                image,
                center_width,
                center_height,
                half_width,
                half_height,
                replace=None)

            is_trial_successfull = True

    return image
//  [python]
def crop_normalized_boxes(
    bboxes: tf.Tensor,
    ori_image_size: tf.Tensor,
    new_image_size: tf.Tensor,
    offset: tf.Tensor,
) -> tf.Tensor:
  """Crop normalized boxes.

  Args:
    bboxes: a float tensor of shape [N, 4] representing normalized box
      coordinates.
    ori_image_size: an int tensor of shape [2] representing the original image
      size.
    new_image_size: an int tensor of shape [2] representing the cropped image
      size.
    offset: an int tensor of shape [2] representing the offset of the crop.

  Returns:
    a tensor of shape [N, 4] representing the new normalized bounding box
    coordinates in the new cropped image.
  """
  new_bboxes = box_ops.denormalize_boxes(bboxes, ori_image_size)
  new_bboxes -= tf.tile(tf.to_float(offset), [2])[None, ...]
  new_bboxes = box_ops.normalize_boxes(new_bboxes, new_image_size)
  return tf.clip_by_value(new_bboxes, 0.0, 1.0)