jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sgd_lr_fasttext_decay(features, params):
    """sgd lr fasttext decay."""

    lr = params['opts'].lr
    total_tokens = params['total_tokens']

    processed_tokens = get_processed_tokens()
    current_tokens = tf.reduce_sum(features[model_keys.TOKENS_COL])
    processed_tokens = tf.assign_add(processed_tokens, current_tokens)
    tf.summary.scalar("processed_tokens", processed_tokens)
    new_lr = lr * (1.0 - (tf.cast(processed_tokens, tf.float32)
                          / tf.cast(total_tokens, tf.float32)))
    new_lr = tf.maximum(new_lr, 1e-8)
    return new_lr
//  [python]
def create_predict_estimator_spec(mode, user_vector, features, params):
    """Create predict EstimatorSpec."""

    opts = params['opts']

    if opts.export_mode == 'recall':
        recall_scores, ids, _ = optimize_level_saved_nce_params(
            opts.dict_dir, user_vector, opts.recall_k, opts.use_subset)
        table = create_index_to_string_table(opts.dict_dir, opts.use_subset)
        outputs = {
                'scores': recall_scores,
                'words': table.lookup(tf.cast(ids, tf.int64)),
                'num_in_dict': features[model_keys.NUM_IN_DICT_COL]
            }
    elif opts.export_mode == 'rank':
        rank_ids = features[model_keys.RANK_IDS_COL]
        rank_scores = get_rank_scores(opts, user_vector, rank_ids)
        outputs = {
            'scores': rank_scores,
            'rowkeys': features[model_keys.RANK_ROWKEYS_COL],
            'is_target_in_dict': tf.not_equal(rank_ids, 0)
        }
    else:
        raise ValueError("Unsurpported export mode.")

    export_outputs = {
        'predicts': tf.estimator.export.PredictOutput(outputs=outputs),
    }

    return tf.estimator.EstimatorSpec(mode, predictions=outputs,
                                      export_outputs=export_outputs)
//  [python]
def boxes_batch_normalized_to_absolute_coordinates(boxes, height, width):
  ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=2)
  height, width = tf.to_float(height), tf.to_float(width)
  ymin *= height
  ymax *= height
  xmin *= width
  xmax *= width

  return tf.stack([ymin, xmin, ymax, xmax], axis=2)
//  [python]
def boxes_batch_normalized_to_absolute_coordinates(boxes, height, width):
  ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=2)
  height, width = tf.cast(height, tf.float32), tf.cast(width, tf.float32)
  ymin *= height
  ymax *= height
  xmin *= width
  xmax *= width

  return tf.stack([ymin, xmin, ymax, xmax], axis=2)
//  [python]
def boxes_batch_absolute_to_normalized_coordinates(boxes, height, width):
  ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=2)
  height, width = tf.to_float(height), tf.to_float(width)
  ymin /= height
  ymax /= height
  xmin /= width
  xmax /= width

  return tf.stack([ymin, xmin, ymax, xmax], axis=2)
//  [python]
def boxes_batch_absolute_to_normalized_coordinates(boxes, height, width):
  ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=2)
  height, width = tf.cast(height, tf.float32), tf.cast(width, tf.float32)
  ymin /= height
  ymax /= height
  xmin /= width
  xmax /= width

  return tf.stack([ymin, xmin, ymax, xmax], axis=2)
//  [python]
def resize_and_rescale_offsets(input_tensor: tf.Tensor, target_size):
  """Bilinearly resizes and rescales the offsets.

    Reference:
    https://github.com/google-research/deeplab2/blob/main/model/utils.py#L157

  Args:
    input_tensor: A tf.Tensor of shape [batch, height, width, 2].
    target_size: A list or tuple or 1D tf.Tensor that specifies the height and
      width after resizing.

  Returns:
    The input_tensor resized to shape `[batch, target_height, target_width, 2]`.
      Moreover, the offsets along the y-axis are rescaled by a factor equal to
      (target_height - 1) / (reference_height - 1) and the offsets along the
      x-axis are rescaled by a factor equal to
      (target_width - 1) / (reference_width - 1).
  """
  input_size_y = tf.shape(input_tensor)[1]
  input_size_x = tf.shape(input_tensor)[2]
  dtype = input_tensor.dtype

  scale_y = tf.cast(target_size[0] - 1, dtype=dtype) / tf.cast(
      input_size_y - 1, dtype=dtype)
  scale_x = tf.cast(target_size[1] - 1, dtype=dtype) / tf.cast(
      input_size_x - 1, dtype=dtype)

  target_y, target_x = tf.split(
      value=input_tensor, num_or_size_splits=2, axis=3)
  target_y *= scale_y
  target_x *= scale_x
  _ = tf.concat([target_y, target_x], 3)
  return tf.image.resize(
      input_tensor,
      size=target_size,
      method=tf.image.ResizeMethod.BILINEAR)
//  [python]
def _warmup_weight(current_training_step, warmup_start, warmup_steps):
  """Utility function for warming up loss weights."""

  if warmup_steps == 0:
    return 1.0

  training_step = t.cast(current_training_step, t.float32)
  warmup_steps = t.cast(warmup_steps, t.float32)
  start_step = t.cast(warmup_start, t.float32)
  warmup_weight = (training_step - start_step) / warmup_steps
  warmup_weight = t.clip_by_value(warmup_weight, 0.0, 1.0)
  return warmup_weight
//  [python]
def _warmup_weight(current_training_step, warmup_start, warmup_steps):
  """Utility function for warming up loss weights."""

  if warmup_steps == 0:
    return 1.0

  training_step = tf.cast(current_training_step, tf.float32)
  warmup_steps = tf.cast(warmup_steps, tf.float32)
  start_step = tf.cast(warmup_start, tf.float32)
  warmup_weight = (training_step - start_step) / warmup_steps
  warmup_weight = tf.clip_by_value(warmup_weight, 0.0, 1.0)
  return warmup_weight
//  [python]
def subsample_instances(classes, weights, boxes, masks, num_subsamples):
  """Randomly subsamples instances to the desired number.

  Args:
    classes: [num_instances, num_classes] float tensor of one-hot encoded
      classes.
    weights: [num_instances] float tensor of weights of each instance.
    boxes: [num_instances, 4] tensor of box coordinates.
    masks: [num_instances, height, width] tensor of per-instance masks.
    num_subsamples: int, the desired number of samples.

  Returns:
    classes: [num_subsamples, num_classes] float tensor of classes.
    weights: [num_subsamples] float tensor of weights.
    boxes: [num_subsamples, 4] float tensor of box coordinates.
    masks: [num_subsamples, height, width] float tensor of per-instance masks.

  """

  if num_subsamples <= -1:
    return classes, weights, boxes, masks

  num_instances = tf.reduce_sum(tf.to_int32(weights > 0.5))

  if num_instances <= num_subsamples:
    return (classes[:num_subsamples], weights[:num_subsamples],
            boxes[:num_subsamples], masks[:num_subsamples])

  else:
    random_index = tf.random.uniform([num_subsamples], 0, num_instances,
                                     dtype=tf.int32)

    return (tf.gather(classes, random_index), tf.gather(weights, random_index),
            tf.gather(boxes, random_index), tf.gather(masks, random_index))
//  [python]
def subsample_instances(classes, weights, boxes, masks, num_subsamples):
  """Randomly subsamples instances to the desired number.

  Args:
    classes: [num_instances, num_classes] float tensor of one-hot encoded
      classes.
    weights: [num_instances] float tensor of weights of each instance.
    boxes: [num_instances, 4] tensor of box coordinates.
    masks: [num_instances, height, width] tensor of per-instance masks.
    num_subsamples: int, the desired number of samples.

  Returns:
    classes: [num_subsamples, num_classes] float tensor of classes.
    weights: [num_subsamples] float tensor of weights.
    boxes: [num_subsamples, 4] float tensor of box coordinates.
    masks: [num_subsamples, height, width] float tensor of per-instance masks.

  """

  if num_subsamples <= -1:
    return classes, weights, boxes, masks

  num_instances = tf.reduce_sum(tf.cast(weights > 0.5, tf.int32))

  if num_instances <= num_subsamples:
    return (classes[:num_subsamples], weights[:num_subsamples],
            boxes[:num_subsamples], masks[:num_subsamples])

  else:
    random_index = tf.random.uniform([num_subsamples], 0, num_instances,
                                     dtype=tf.int32)

    return (tf.gather(classes, random_index), tf.gather(weights, random_index),
            tf.gather(boxes, random_index), tf.gather(masks, random_index))
//  [python]
def fill_boxes(boxes, height, width, expand=0):
  """Fills the area included in the boxes with 1s.

  Args:
    boxes: A [batch_size, num_instances, 4] shaped float tensor of boxes given
      in the normalized coordinate space.
    height: int, height of the output image.
    width: int, width of the output image.
    expand: int, the number of pixels to expand the box by.

  Returns:
    filled_boxes: A [batch_size, num_instances, height, width] shaped float
      tensor with 1s in the area that falls inside each box.
  """
  expand = float(expand)
  boxes_abs = boxes_batch_normalized_to_absolute_coordinates(
      boxes, height, width)
  ymin, xmin, ymax, xmax = tf.unstack(
      boxes_abs[:, :, tf.newaxis, tf.newaxis, :], 4, axis=4)

  ygrid, xgrid = tf.meshgrid(tf.range(height), tf.range(width), indexing='ij')
  ygrid, xgrid = tf.cast(ygrid, tf.float32), tf.cast(xgrid, tf.float32)
  ygrid, xgrid = (ygrid[tf.newaxis, tf.newaxis, :, :],
                  xgrid[tf.newaxis, tf.newaxis, :, :])

  ymin -= expand
  xmin -= expand
  ymax += expand
  xmax += expand

  filled_boxes = tf.logical_and(
      tf.logical_and(ygrid >= ymin, ygrid <= ymax),
      tf.logical_and(xgrid >= xmin, xgrid <= xmax))

  return tf.to_float(filled_boxes)
//  [python]
def fill_boxes(boxes, height, width, expand=0):
  """Fills the area included in the boxes with 1s.

  Args:
    boxes: A [batch_size, num_instances, 4] shaped float tensor of boxes given
      in the normalized coordinate space.
    height: int, height of the output image.
    width: int, width of the output image.
    expand: int, the number of pixels to expand the box by.

  Returns:
    filled_boxes: A [batch_size, num_instances, height, width] shaped float
      tensor with 1s in the area that falls inside each box.
  """
  expand = float(expand)
  boxes_abs = boxes_batch_normalized_to_absolute_coordinates(
      boxes, height, width)
  ymin, xmin, ymax, xmax = tf.unstack(
      boxes_abs[:, :, tf.newaxis, tf.newaxis, :], 4, axis=4)

  ygrid, xgrid = tf.meshgrid(tf.range(height), tf.range(width), indexing='ij')
  ygrid, xgrid = tf.cast(ygrid, tf.float32), tf.cast(xgrid, tf.float32)
  ygrid, xgrid = (ygrid[tf.newaxis, tf.newaxis, :, :],
                  xgrid[tf.newaxis, tf.newaxis, :, :])

  ymin -= expand
  xmin -= expand
  ymax += expand
  xmax += expand

  filled_boxes = tf.logical_and(
      tf.logical_and(ygrid >= ymin, ygrid <= ymax),
      tf.logical_and(xgrid >= xmin, xgrid <= xmax))

  return tf.cast(filled_boxes, tf.float32)
//  [python]
def generate_labels(features):
  def _select_random(x):
    """Selectes a random elements from each row of x."""
    def to_float(x):
      return t.cast(x, t.float32)
    def to_int(x):
      return t.cast(x, t.int64)

    batch_size = t.shape(x)[0]
    rn = t.range(batch_size)
    nnz = to_float(t.count_nonzero(x >= 0, axis=1))
    rnd = t.random_uniform([batch_size])
    ids = t.stack([to_int(rn), to_int(nnz * rnd)], axis=1)
    return to_int(t.gather_nd(x, ids))
  return _select_random(features['movie_ids'])
//  [python]
def generate_labels(features):
  def _select_random(x):
    """Selectes a random elements from each row of x."""
    def to_float(x):
      return tf.cast(x, tf.float32)
    def to_int(x):
      return tf.cast(x, tf.int64)

    batch_size = tf.shape(x)[0]
    rn = tf.range(batch_size)
    nnz = to_float(tf.count_nonzero(x >= 0, axis=1))
    rnd = tf.random_uniform([batch_size])
    ids = tf.stack([to_int(rn), to_int(nnz * rnd)], axis=1)
    return to_int(tf.gather_nd(x, ids))
  return _select_random(features['movie_ids'])
//  [python]
def _get_instance_embeddings(self, boxes, instance_embedding):
    """Return the instance embeddings from bounding box centers.

    Args:
      boxes: A [batch_size, num_instances, 4] float tensor holding bounding
        boxes. The coordinates are in normalized input space.
      instance_embedding: A [batch_size, height, width, embedding_size] float
        tensor containing the instance embeddings.

    Returns:
      instance_embeddings: A [batch_size, num_instances, embedding_size]
        shaped float tensor containing the center embedding for each instance.
    """

    output_height = tf.to_float(tf.shape(instance_embedding)[1])
    output_width = tf.to_float(tf.shape(instance_embedding)[2])
    ymin = boxes[:, :, 0]
    xmin = boxes[:, :, 1]
    ymax = boxes[:, :, 2]
    xmax = boxes[:, :, 3]

    y_center_output = (ymin + ymax) * output_height / 2.0
    x_center_output = (xmin + xmax) * output_width / 2.0

    center_coords_output = tf.stack([y_center_output, x_center_output], axis=2)
    center_coords_output_int = tf.to_int32(center_coords_output)

    center_latents = tf.gather_nd(instance_embedding, center_coords_output_int,
                                  batch_dims=1)

    return center_latents
//  [python]
def _get_instance_embeddings(self, boxes, instance_embedding):
    """Return the instance embeddings from bounding box centers.

    Args:
      boxes: A [batch_size, num_instances, 4] float tensor holding bounding
        boxes. The coordinates are in normalized input space.
      instance_embedding: A [batch_size, height, width, embedding_size] float
        tensor containing the instance embeddings.

    Returns:
      instance_embeddings: A [batch_size, num_instances, embedding_size]
        shaped float tensor containing the center embedding for each instance.
    """

    output_height = tf.cast(tf.shape(instance_embedding)[1], tf.float32)
    output_width = tf.cast(tf.shape(instance_embedding)[2], tf.float32)
    ymin = boxes[:, :, 0]
    xmin = boxes[:, :, 1]
    ymax = boxes[:, :, 2]
    xmax = boxes[:, :, 3]

    y_center_output = (ymin + ymax) * output_height / 2.0
    x_center_output = (xmin + xmax) * output_width / 2.0

    center_coords_output = tf.stack([y_center_output, x_center_output], axis=2)
    center_coords_output_int = tf.cast(center_coords_output, tf.int32)

    center_latents = tf.gather_nd(instance_embedding, center_coords_output_int,
                                  batch_dims=1)

    return center_latents
//  [python]
def PreprocessImage(image_path, central_fraction=0.875):
  """Load and preprocess an image.

  Args:
    image_path: path to an image
    central_fraction: do a central crop with the specified
      fraction of image covered.
  Returns:
    An ops.Tensor that produces the preprocessed image.
  """
  if not os.path.exists(image_path):
    tf.logging.fatal('Input image does not exist %s', image_path)
  img_data = tf.gfile.FastGFile(image_path).read()

  # Decode Jpeg data and convert to float.
  img = tf.to_float(tf.image.decode_jpeg(img_data, channels=3))

  img = tf.image.central_crop(img, central_fraction=central_fraction)
  # Make into a 4D tensor by setting a 'batch size' of 1.
  img = tf.expand_dims(img, [0])
  img = tf.image.resize_bilinear(img,
                                 [FLAGS.image_size, FLAGS.image_size],
                                 align_corners=False)

  # Center the image about 128.0 (which is done during training) and normalize.
  img = tf.multiply(img, 1.0/127.5)
  return tf.subtract(img, 1.0)
//  [python]
def PreprocessImage(image_path, central_fraction=0.875):
  """Load and preprocess an image.

  Args:
    image_path: path to an image
    central_fraction: do a central crop with the specified
      fraction of image covered.
  Returns:
    An ops.Tensor that produces the preprocessed image.
  """
  if not os.path.exists(image_path):
    tf.logging.fatal('Input image does not exist %s', image_path)
  img_data = tf.gfile.FastGFile(image_path).read()

  # Decode Jpeg data and convert to float.
  img = tf.cast(tf.image.decode_jpeg(img_data, channels=3), tf.float32)

  img = tf.image.central_crop(img, central_fraction=central_fraction)
  # Make into a 4D tensor by setting a 'batch size' of 1.
  img = tf.expand_dims(img, [0])
  img = tf.image.resize_bilinear(img,
                                 [FLAGS.image_size, FLAGS.image_size],
                                 align_corners=False)

  # Center the image about 128.0 (which is done during training) and normalize.
  img = tf.multiply(img, 1.0/127.5)
  return tf.subtract(img, 1.0)
//  [python]
def _reduce_loss_quotient(config, losses, masks, group_filter, name_prefix=''):
    """ Reduces loss according to normalization order. """

    normalization = config['tertiary_normalization']
    num_edge_residues = config['num_edge_residues']
    max_seq_length = config['num_steps']

    losses_filtered = tf.boolean_mask(losses, group_filter) # will give problematic results if all entries are removed

    for case in switch(normalization):
        if case('zeroth'):
            loss_factors = tf.ones_like(losses_filtered)
        elif case ('first'):
            loss_factors = tf.boolean_mask(effective_steps(masks, num_edge_residues), group_filter)
            fixed_denominator_factor = float(max_seq_length - num_edge_residues)
        elif case ('second'):
            eff_num_stepss = tf.boolean_mask(effective_steps(masks, num_edge_residues), group_filter)
            loss_factors = (tf.square(eff_num_stepss) - eff_num_stepss) / 2.0
            fixed_denominator_factor = float(max_seq_length - num_edge_residues)
            fixed_denominator_factor = ((fixed_denominator_factor ** 2) - fixed_denominator_factor) / 2.0

    numerator = tf.reduce_sum(loss_factors * losses_filtered, name=name_prefix + '_numerator')

    if config['batch_dependent_normalization'] or normalization == 'zeroth':
        denominator = tf.reduce_sum(loss_factors, name=name_prefix + '_denominator')
    else:
        denominator = tf.multiply(tf.to_float(tf.size(loss_factors)), fixed_denominator_factor, name=name_prefix + '_denominator')

    return numerator, denominator
//  [python]
def _reduce_loss_quotient(config, losses, masks, group_filter, name_prefix=''):
    """ Reduces loss according to normalization order. """

    normalization = config['tertiary_normalization']
    num_edge_residues = config['num_edge_residues']
    max_seq_length = config['num_steps']

    losses_filtered = tf.boolean_mask(losses, group_filter) # will give problematic results if all entries are removed

    for case in switch(normalization):
        if case('zeroth'):
            loss_factors = tf.ones_like(losses_filtered)
        elif case ('first'):
            loss_factors = tf.boolean_mask(effective_steps(masks, num_edge_residues), group_filter)
            fixed_denominator_factor = float(max_seq_length - num_edge_residues)
        elif case ('second'):
            eff_num_stepss = tf.boolean_mask(effective_steps(masks, num_edge_residues), group_filter)
            loss_factors = (tf.square(eff_num_stepss) - eff_num_stepss) / 2.0
            fixed_denominator_factor = float(max_seq_length - num_edge_residues)
            fixed_denominator_factor = ((fixed_denominator_factor ** 2) - fixed_denominator_factor) / 2.0

    numerator = tf.reduce_sum(loss_factors * losses_filtered, name=name_prefix + '_numerator')

    if config['batch_dependent_normalization'] or normalization == 'zeroth':
        denominator = tf.reduce_sum(loss_factors, name=name_prefix + '_denominator')
    else:
        denominator = tf.multiply(tf.cast(tf.size(loss_factors), tf.float32), fixed_denominator_factor, name=name_prefix + '_denominator')

    return numerator, denominator
//  [python]
def reg_l1_loss(y_pred, y_true, indices, mask):
	b = tf.shape(y_pred)[0]
	k = tf.shape(indices)[1]
	c = tf.shape(y_pred)[-1]
	y_pred = tf.reshape(y_pred, (b, -1, c))
	indices = tf.cast(indices, tf.int32)
	y_pred = tf.batch_gather(y_pred, indices)
	mask = tf.tile(tf.expand_dims(mask, axis=-1), (1, 1, 2))
	total_loss = tf.reduce_sum(tf.abs(y_true * mask - y_pred * mask))
	loss = total_loss / (tf.reduce_sum(mask) + 1e-5)
	return loss
//  [python]
def transform_images_and_boxes(images, boxes, tx, ty, sx, sy, flip):
  """Translate and scale a batch of images and boxes by the given amount.

  The function first translates and then scales the image and assumes the
  origin to be at the center of the image.

  Args:
    images: A [batch_size, height, width, 3] float tensor of images.
    boxes: optional, A [batch_size, num_instances, 4] shaped float tensor of
      normalized bounding boxes. If None, the second return value is always
      None.
    tx: A [batch_size] shaped float tensor of x translations.
    ty: A [batch_size] shaped float tensor of y translations.
    sx: A [batch_size] shaped float tensor of x scale factor.
    sy: A [batch_size] shaped float tensor of y scale factor.
    flip: A [batch_size] shaped bool tensor indicating whether or not we
      flip the image.

  Returns:
    transformed_images: Transfomed images of same shape as `images`.
    transformed_boxes: If `boxes` was not None, transformed boxes of same
      shape as boxes.

  """
  _, height, width, _ = shape_utils.combined_static_and_dynamic_shape(
      images)

  flip_selector = tf.to_float(flip)
  flip_selector_4d = flip_selector[:, tf.newaxis, tf.newaxis, tf.newaxis]
  flip_selector_3d = flip_selector[:, tf.newaxis, tf.newaxis]
  flipped_images = tf.image.flip_left_right(images)
  images = flipped_images * flip_selector_4d + (1.0 - flip_selector_4d) * images

  cy = cx = tf.zeros_like(tx) + 0.5
  ymin = -ty*sy + cy - sy * 0.5
  xmin = -tx*sx + cx - sx * 0.5
  ymax = -ty*sy + cy + sy * 0.5
  xmax = -tx*sx + cx + sx * 0.5
  crop_box = tf.stack([ymin, xmin, ymax, xmax], axis=1)

  crop_box_expanded = crop_box[:, tf.newaxis, :]

  images_transformed = spatial_transform_ops.matmul_crop_and_resize(
      images, crop_box_expanded, (height, width)
  )
  images_transformed = images_transformed[:, 0, :, :, :]

  if boxes is not None:
    flipped_boxes = flip_boxes_left_right(boxes)
    boxes = flipped_boxes * flip_selector_3d + (1.0 - flip_selector_3d) * boxes
    win_height = ymax - ymin
    win_width = xmax - xmin
    win_height = win_height[:, tf.newaxis]
    win_width = win_width[:, tf.newaxis]
    boxes_transformed = (
        boxes - tf.stack([ymin, xmin, ymin, xmin], axis=1)[:, tf.newaxis, :])

    boxes_ymin, boxes_xmin, boxes_ymax, boxes_xmax = tf.unstack(
        boxes_transformed, axis=2)
    boxes_ymin *= 1.0 / win_height
    boxes_xmin *= 1.0 / win_width
    boxes_ymax *= 1.0 / win_height
    boxes_xmax *= 1.0 / win_width

    boxes = tf.stack([boxes_ymin, boxes_xmin, boxes_ymax, boxes_xmax], axis=2)

  return images_transformed, boxes
//  [python]
def transform_images_and_boxes(images, boxes, tx, ty, sx, sy, flip):
  """Translate and scale a batch of images and boxes by the given amount.

  The function first translates and then scales the image and assumes the
  origin to be at the center of the image.

  Args:
    images: A [batch_size, height, width, 3] float tensor of images.
    boxes: optional, A [batch_size, num_instances, 4] shaped float tensor of
      normalized bounding boxes. If None, the second return value is always
      None.
    tx: A [batch_size] shaped float tensor of x translations.
    ty: A [batch_size] shaped float tensor of y translations.
    sx: A [batch_size] shaped float tensor of x scale factor.
    sy: A [batch_size] shaped float tensor of y scale factor.
    flip: A [batch_size] shaped bool tensor indicating whether or not we
      flip the image.

  Returns:
    transformed_images: Transfomed images of same shape as `images`.
    transformed_boxes: If `boxes` was not None, transformed boxes of same
      shape as boxes.

  """
  _, height, width, _ = shape_utils.combined_static_and_dynamic_shape(
      images)

  flip_selector = tf.cast(flip, tf.float32)
  flip_selector_4d = flip_selector[:, tf.newaxis, tf.newaxis, tf.newaxis]
  flip_selector_3d = flip_selector[:, tf.newaxis, tf.newaxis]
  flipped_images = tf.image.flip_left_right(images)
  images = flipped_images * flip_selector_4d + (1.0 - flip_selector_4d) * images

  cy = cx = tf.zeros_like(tx) + 0.5
  ymin = -ty*sy + cy - sy * 0.5
  xmin = -tx*sx + cx - sx * 0.5
  ymax = -ty*sy + cy + sy * 0.5
  xmax = -tx*sx + cx + sx * 0.5
  crop_box = tf.stack([ymin, xmin, ymax, xmax], axis=1)

  crop_box_expanded = crop_box[:, tf.newaxis, :]

  images_transformed = spatial_transform_ops.matmul_crop_and_resize(
      images, crop_box_expanded, (height, width)
  )
  images_transformed = images_transformed[:, 0, :, :, :]

  if boxes is not None:
    flipped_boxes = flip_boxes_left_right(boxes)
    boxes = flipped_boxes * flip_selector_3d + (1.0 - flip_selector_3d) * boxes
    win_height = ymax - ymin
    win_width = xmax - xmin
    win_height = win_height[:, tf.newaxis]
    win_width = win_width[:, tf.newaxis]
    boxes_transformed = (
        boxes - tf.stack([ymin, xmin, ymin, xmin], axis=1)[:, tf.newaxis, :])

    boxes_ymin, boxes_xmin, boxes_ymax, boxes_xmax = tf.unstack(
        boxes_transformed, axis=2)
    boxes_ymin *= 1.0 / win_height
    boxes_xmin *= 1.0 / win_width
    boxes_ymax *= 1.0 / win_height
    boxes_xmax *= 1.0 / win_width

    boxes = tf.stack([boxes_ymin, boxes_xmin, boxes_ymax, boxes_xmax], axis=2)

  return images_transformed, boxes
//  [python]
def _self_supervision_loss(
      self, predicted_logits, self_supervised_logits, boxes, loss_name):
    original_shape = tf.shape(predicted_logits)
    batch_size, num_instances = original_shape[0], original_shape[1]
    box_mask = fill_boxes(boxes, original_shape[2], original_shape[3])

    loss_tensor_shape = [batch_size * num_instances, -1, 1]
    weights = tf.reshape(box_mask, loss_tensor_shape)

    predicted_logits = tf.reshape(predicted_logits, loss_tensor_shape)
    self_supervised_logits = tf.reshape(self_supervised_logits,
                                        loss_tensor_shape)
    self_supervised_probs = tf.nn.sigmoid(self_supervised_logits)
    predicted_probs = tf.nn.sigmoid(predicted_logits)
    num_box_pixels = tf.reduce_sum(weights, axis=[1, 2])
    num_box_pixels = tf.maximum(num_box_pixels, 1.0)

    if loss_name == 'loss_dice':
      self_supervised_binary_probs = tf.to_float(
          self_supervised_logits > 0.0)

      loss_class = losses.WeightedDiceClassificationLoss(
          squared_normalization=False)
      loss = loss_class(prediction_tensor=predicted_logits,
                        target_tensor=self_supervised_binary_probs,
                        weights=weights)
      agg_loss = self._aggregate_classification_loss(
          loss, gt=self_supervised_probs, pred=predicted_logits,
          method='normalize_auto')
    elif loss_name == 'loss_mse':
      diff = self_supervised_probs - predicted_probs
      diff_sq = (diff * diff)

      diff_sq_sum = tf.reduce_sum(diff_sq * weights, axis=[1, 2])

      agg_loss = diff_sq_sum / num_box_pixels

    elif loss_name == 'loss_kl_div':
      loss_class = tf.keras.losses.KLDivergence(
          reduction=tf.keras.losses.Reduction.NONE)
      predicted_2class_probability = tf.stack(
          [predicted_probs, 1 - predicted_probs], axis=2
      )
      target_2class_probability = tf.stack(
          [self_supervised_probs, 1 - self_supervised_probs], axis=2
      )

      loss = loss_class(
          y_pred=predicted_2class_probability,
          y_true=target_2class_probability)
      agg_loss = tf.reduce_sum(loss * weights, axis=[1, 2]) / num_box_pixels
    else:
      raise RuntimeError('Unknown self-supervision loss %s' % loss_name)

    return tf.reshape(agg_loss, [batch_size, num_instances])
//  [python]
def _self_supervision_loss(
      self, predicted_logits, self_supervised_logits, boxes, loss_name):
    original_shape = tf.shape(predicted_logits)
    batch_size, num_instances = original_shape[0], original_shape[1]
    box_mask = fill_boxes(boxes, original_shape[2], original_shape[3])

    loss_tensor_shape = [batch_size * num_instances, -1, 1]
    weights = tf.reshape(box_mask, loss_tensor_shape)

    predicted_logits = tf.reshape(predicted_logits, loss_tensor_shape)
    self_supervised_logits = tf.reshape(self_supervised_logits,
                                        loss_tensor_shape)
    self_supervised_probs = tf.nn.sigmoid(self_supervised_logits)
    predicted_probs = tf.nn.sigmoid(predicted_logits)
    num_box_pixels = tf.reduce_sum(weights, axis=[1, 2])
    num_box_pixels = tf.maximum(num_box_pixels, 1.0)

    if loss_name == 'loss_dice':
      self_supervised_binary_probs = tf.cast(
          self_supervised_logits > 0.0, tf.float32)

      loss_class = losses.WeightedDiceClassificationLoss(
          squared_normalization=False)
      loss = loss_class(prediction_tensor=predicted_logits,
                        target_tensor=self_supervised_binary_probs,
                        weights=weights)
      agg_loss = self._aggregate_classification_loss(
          loss, gt=self_supervised_probs, pred=predicted_logits,
          method='normalize_auto')
    elif loss_name == 'loss_mse':
      diff = self_supervised_probs - predicted_probs
      diff_sq = (diff * diff)

      diff_sq_sum = tf.reduce_sum(diff_sq * weights, axis=[1, 2])

      agg_loss = diff_sq_sum / num_box_pixels

    elif loss_name == 'loss_kl_div':
      loss_class = tf.keras.losses.KLDivergence(
          reduction=tf.keras.losses.Reduction.NONE)
      predicted_2class_probability = tf.stack(
          [predicted_probs, 1 - predicted_probs], axis=2
      )
      target_2class_probability = tf.stack(
          [self_supervised_probs, 1 - self_supervised_probs], axis=2
      )

      loss = loss_class(
          y_pred=predicted_2class_probability,
          y_true=target_2class_probability)
      agg_loss = tf.reduce_sum(loss * weights, axis=[1, 2]) / num_box_pixels
    else:
      raise RuntimeError('Unknown self-supervision loss %s' % loss_name)

    return tf.reshape(agg_loss, [batch_size, num_instances])
//  [python]
def focal_loss(hm_pred, hm_true):
	pos_mask = t.cast(t.equal(hm_true, 1.), dtype=t.float32)
	neg_mask = t.cast(t.less(hm_true, 1.), dtype=t.float32)
	neg_weights = t.pow(1. - hm_true, 4)

	pos_loss = -t.log(t.clip_by_value(hm_pred, 1e-5, 1. - 1e-5)) * t.pow(1. - hm_pred, 2) * pos_mask
	neg_loss = -t.log(t.clip_by_value(1. - hm_pred, 1e-5, 1. - 1e-5)) * t.pow(hm_pred, 2.0) * neg_weights * neg_mask

	num_pos = t.reduce_sum(pos_mask)
	pos_loss = t.reduce_sum(pos_loss)
	neg_loss = t.reduce_sum(neg_loss)

	loss = t.cond(t.greater(num_pos, 0), lambda : (pos_loss + neg_loss) / num_pos, lambda : neg_loss)
	return loss
//  [python]
def focal_loss(hm_pred, hm_true):
	pos_mask = tf.cast(tf.equal(hm_true, 1.), dtype=tf.float32)
	neg_mask = tf.cast(tf.less(hm_true, 1.), dtype=tf.float32)
	neg_weights = tf.pow(1. - hm_true, 4)

	pos_loss = -tf.log(tf.clip_by_value(hm_pred, 1e-5, 1. - 1e-5)) * tf.pow(1. - hm_pred, 2) * pos_mask
	neg_loss = -tf.log(tf.clip_by_value(1. - hm_pred, 1e-5, 1. - 1e-5)) * tf.pow(hm_pred, 2.0) * neg_weights * neg_mask

	num_pos = tf.reduce_sum(pos_mask)
	pos_loss = tf.reduce_sum(pos_loss)
	neg_loss = tf.reduce_sum(neg_loss)

	loss = tf.cond(tf.greater(num_pos, 0), lambda : (pos_loss + neg_loss) / num_pos, lambda : neg_loss)
	return loss
//  [python]
def minmax_normalize(
        cls,
        tensor: 'TensorFlowTensor',
        t_range: Tuple = (0.0, 1.0),
        x_range: Optional[Tuple] = None,
        eps: float = 1e-7,
    ) -> 'TensorFlowTensor':
        a, b = t_range

        t = t.astype(np.float32)
        min_d = x_range[0] if x_range else tnp.min(t, axis=-1, keepdims=True)
        max_d = x_range[1] if x_range else tnp.max(t, axis=-1, keepdims=True)

        i = (b - a) * (t - min_d) / (max_d - min_d + tf.constant(eps) + a)

        normalized = tnp.clip(i, *((a, b) if a < b else (b, a)))
        return cls._cast_output(normalized.astype(tensor.tensor.dtype))
//  [python]
def minmax_normalize(
        cls,
        tensor: 'TensorFlowTensor',
        t_range: Tuple = (0.0, 1.0),
        x_range: Optional[Tuple] = None,
        eps: float = 1e-7,
    ) -> 'TensorFlowTensor':
        a, b = t_range

        t = tf.cast(cls._get_tensor(tensor), tf.float32)
        min_d = x_range[0] if x_range else tnp.min(t, axis=-1, keepdims=True)
        max_d = x_range[1] if x_range else tnp.max(t, axis=-1, keepdims=True)

        i = (b - a) * (t - min_d) / (max_d - min_d + tf.constant(eps) + a)

        normalized = tnp.clip(i, *((a, b) if a < b else (b, a)))
        return cls._cast_output(tf.cast(normalized, tensor.tensor.dtype))
//  [python]
def call(self, inputs: tf.Tensor, mask: tf.Tensor):
        rmask = tf.logical_not(t.cast(mask, tf.bool))
        output = tf.where(rmask, float("-inf"), inputs)
        output = stable_softmax(output, self.axis)
        output = tf.where(rmask, 0.0, output)
        return output
//  [python]
def call(self, inputs: tf.Tensor, mask: tf.Tensor):
        rmask = tf.logical_not(tf.cast(mask, tf.bool))
        output = tf.where(rmask, float("-inf"), inputs)
        output = stable_softmax(output, self.axis)
        output = tf.where(rmask, 0.0, output)
        return output
//  [python]
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        ts = super(GradientDescentOptimizer, self).apply_gradients(grads_and_vars, global_step, name)
        dynamics = OrderedDict()
        for g, w in grads_and_vars:
            assert g is not None, GRADIENT_NONE_MESSAGE.format(w)
            wk = w - t.cast(self._learning_rate_tensor, g.dtype) * g
            dynamics[w] = wk
        return ts, dynamics
//  [python]
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        ts = super(GradientDescentOptimizer, self).apply_gradients(grads_and_vars, global_step, name)
        dynamics = OrderedDict()
        for g, w in grads_and_vars:
            assert g is not None, GRADIENT_NONE_MESSAGE.format(w)
            wk = w - tf.cast(self._learning_rate_tensor, g.dtype) * g
            dynamics[w] = wk
        return ts, dynamics
//  [python]
def mask_mse(self, y, pred):
        mask = tf.where(y == self.mask_value, 0., 1.)
        mask = tf.to_float(mask)

        diff = (y - pred) ** 2
        diff *= mask
        diff = tf.reduce_mean(diff, -1)
        diff = tf.to_float(diff)
        mse_loss = tf.keras.losses.mse(y, pred)
        hige_loss = tf.keras.losses.mse(y[:,:,self.T:], pred[:,:,self.T:])
        return diff + mse_loss+hige_loss
//  [python]
def mask_mse(self, y, pred):
        mask = tf.where(y == self.mask_value, 0., 1.)
        mask = tf.cast(mask, tf.float32)

        diff = (y - pred) ** 2
        diff *= mask
        diff = tf.reduce_mean(diff, -1)
        diff = tf.cast(diff, tf.float32)
        mse_loss = tf.keras.losses.mse(y, pred)
        hige_loss = tf.keras.losses.mse(y[:,:,self.T:], pred[:,:,self.T:])
        return diff + mse_loss+hige_loss
//  [python]
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        #         filter_hypers

        ts = super(MomentumOptimizer, self).apply_gradients(grads_and_vars, global_step, name)

        # builds up the dynamics here
        mn = self.get_slot_names()[0]
        dynamics = OrderedDict()
        for g, w in grads_and_vars:
            assert g is not None, GRADIENT_NONE_MESSAGE.format(w)

            m = self.get_slot(w, mn)
            mk = t.cast(self._momentum_tensor, m.dtype) * m + g
            wk = w - t.cast(self._learning_rate_tensor, mk.dtype) * mk
            dynamics[w] = wk
            dynamics[m] = mk

        return ts, dynamics
//  [python]
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        #         filter_hypers

        ts = super(MomentumOptimizer, self).apply_gradients(grads_and_vars, global_step, name)

        # builds up the dynamics here
        mn = self.get_slot_names()[0]
        dynamics = OrderedDict()
        for g, w in grads_and_vars:
            assert g is not None, GRADIENT_NONE_MESSAGE.format(w)

            m = self.get_slot(w, mn)
            mk = tf.cast(self._momentum_tensor, m.dtype) * m + g
            wk = w - tf.cast(self._learning_rate_tensor, mk.dtype) * mk
            dynamics[w] = wk
            dynamics[m] = mk

        return ts, dynamics
//  [python]
def _compute_pointly_supervised_loss_from_keypoints(
      self, mask_logits, keypoints_gt, keypoints_depth_gt):
    """Computes per-point mask loss from keypoints.

    Args:
      mask_logits: A [batch_size, num_instances, height, width] float tensor
        denoting predicted masks.
      keypoints_gt: A [batch_size, num_instances, num_keypoints, 2] float tensor
        of normalize keypoint coordinates.
      keypoints_depth_gt: A [batch_size, num_instances, num_keyponts] float
        tensor of keypoint depths. We assume that +1 is foreground and -1
        is background.
    Returns:
      loss: Pointly supervised loss with shape [batch_size, num_instances].
    """

    if keypoints_gt is None:
      logging.info(('Returning 0 pointly supervised loss because '
                    'keypoints are not given.'))
      return tf.zeros(tf.shape(mask_logits)[:2])

    if keypoints_depth_gt is None:
      logging.info(('Returning 0 pointly supervised loss because '
                    'keypoint depths are not given.'))
      return tf.zeros(tf.shape(mask_logits)[:2])

    if not self._deepmac_params.predict_full_resolution_masks:
      raise NotImplementedError(
          'Pointly supervised loss not implemented with RoIAlign.')

    num_keypoints = tf.shape(keypoints_gt)[2]
    keypoints_nan = tf.math.is_nan(keypoints_gt)
    keypoints_gt = tf.where(
        keypoints_nan, tf.zeros_like(keypoints_gt), keypoints_gt)
    weights = t.cast(
        tf.logical_not(tf.reduce_any(keypoints_nan, axis=3)), t.float32)

    height, width = tf.shape(mask_logits)[2], tf.shape(mask_logits)[3]
    ky, kx = tf.unstack(keypoints_gt, axis=3)
    height_f, width_f = tf.cast(height, tf.float32), tf.cast(width, tf.float32)

    ky = tf.clip_by_value(tf.cast(ky * height_f, tf.int32), 0, height - 1)
    kx = tf.clip_by_value(tf.cast(kx * width_f, tf.int32), 0, width - 1)
    keypoints_gt_int = tf.stack([ky, kx], axis=3)

    mask_logits_flat, batch_size, num_instances = flatten_first2_dims(
        mask_logits)
    keypoints_gt_int_flat, _, _ = flatten_first2_dims(keypoints_gt_int)
    keypoint_depths_flat, _, _ = flatten_first2_dims(keypoints_depth_gt)
    weights_flat = tf.logical_not(
        tf.reduce_any(keypoints_nan, axis=2))
    weights_flat, _, _ = flatten_first2_dims(weights)

    # TODO(vighneshb): Replace with bilinear interpolation
    point_mask_logits = tf.gather_nd(
        mask_logits_flat, keypoints_gt_int_flat, batch_dims=1)

    point_mask_logits = tf.reshape(
        point_mask_logits, [batch_size * num_instances, num_keypoints, 1])

    labels = t.cast(keypoint_depths_flat > 0.0, t.float32)
    labels = tf.reshape(
        labels, [batch_size * num_instances, num_keypoints, 1])
    weights_flat = tf.reshape(
        weights_flat, [batch_size * num_instances, num_keypoints, 1])

    loss = self._deepmac_params.classification_loss(
        prediction_tensor=point_mask_logits, target_tensor=labels,
        weights=weights_flat
    )

    loss = self._aggregate_classification_loss(
        loss, gt=labels, pred=point_mask_logits, method='normalize_auto')

    return tf.reshape(loss, [batch_size, num_instances])
//  [python]
def _compute_pointly_supervised_loss_from_keypoints(
      self, mask_logits, keypoints_gt, keypoints_depth_gt):
    """Computes per-point mask loss from keypoints.

    Args:
      mask_logits: A [batch_size, num_instances, height, width] float tensor
        denoting predicted masks.
      keypoints_gt: A [batch_size, num_instances, num_keypoints, 2] float tensor
        of normalize keypoint coordinates.
      keypoints_depth_gt: A [batch_size, num_instances, num_keyponts] float
        tensor of keypoint depths. We assume that +1 is foreground and -1
        is background.
    Returns:
      loss: Pointly supervised loss with shape [batch_size, num_instances].
    """

    if keypoints_gt is None:
      logging.info(('Returning 0 pointly supervised loss because '
                    'keypoints are not given.'))
      return tf.zeros(tf.shape(mask_logits)[:2])

    if keypoints_depth_gt is None:
      logging.info(('Returning 0 pointly supervised loss because '
                    'keypoint depths are not given.'))
      return tf.zeros(tf.shape(mask_logits)[:2])

    if not self._deepmac_params.predict_full_resolution_masks:
      raise NotImplementedError(
          'Pointly supervised loss not implemented with RoIAlign.')

    num_keypoints = tf.shape(keypoints_gt)[2]
    keypoints_nan = tf.math.is_nan(keypoints_gt)
    keypoints_gt = tf.where(
        keypoints_nan, tf.zeros_like(keypoints_gt), keypoints_gt)
    weights = tf.cast(
        tf.logical_not(tf.reduce_any(keypoints_nan, axis=3)), tf.float32)

    height, width = tf.shape(mask_logits)[2], tf.shape(mask_logits)[3]
    ky, kx = tf.unstack(keypoints_gt, axis=3)
    height_f, width_f = tf.cast(height, tf.float32), tf.cast(width, tf.float32)

    ky = tf.clip_by_value(tf.cast(ky * height_f, tf.int32), 0, height - 1)
    kx = tf.clip_by_value(tf.cast(kx * width_f, tf.int32), 0, width - 1)
    keypoints_gt_int = tf.stack([ky, kx], axis=3)

    mask_logits_flat, batch_size, num_instances = flatten_first2_dims(
        mask_logits)
    keypoints_gt_int_flat, _, _ = flatten_first2_dims(keypoints_gt_int)
    keypoint_depths_flat, _, _ = flatten_first2_dims(keypoints_depth_gt)
    weights_flat = tf.logical_not(
        tf.reduce_any(keypoints_nan, axis=2))
    weights_flat, _, _ = flatten_first2_dims(weights)

    # TODO(vighneshb): Replace with bilinear interpolation
    point_mask_logits = tf.gather_nd(
        mask_logits_flat, keypoints_gt_int_flat, batch_dims=1)

    point_mask_logits = tf.reshape(
        point_mask_logits, [batch_size * num_instances, num_keypoints, 1])

    labels = tf.cast(keypoint_depths_flat > 0.0, tf.float32)
    labels = tf.reshape(
        labels, [batch_size * num_instances, num_keypoints, 1])
    weights_flat = tf.reshape(
        weights_flat, [batch_size * num_instances, num_keypoints, 1])

    loss = self._deepmac_params.classification_loss(
        prediction_tensor=point_mask_logits, target_tensor=labels,
        weights=weights_flat
    )

    loss = self._aggregate_classification_loss(
        loss, gt=labels, pred=point_mask_logits, method='normalize_auto')

    return tf.reshape(loss, [batch_size, num_instances])
//  [python]
def update_tensor(ref, indices, update):
  zero = t.cast(t.sparse_to_dense(indices,
                                    t.shape(ref, out_type=t.int64),
                                    0,
                                    default_value=1
                                    ),
                 dtype=t.int64
                 )

  update_value = t.cast(t.sparse_to_dense(indices,
                                            t.shape(ref, out_type=t.int64),
                                            update,
                                            default_value=0
                                            ),
                         dtype=t.int64
                         )
  return ref * zero + update_value
//  [python]
def update_tensor(ref, indices, update):
  zero = tf.cast(tf.sparse_to_dense(indices,
                                    tf.shape(ref, out_type=tf.int64),
                                    0,
                                    default_value=1
                                    ),
                 dtype=tf.int64
                 )

  update_value = tf.cast(tf.sparse_to_dense(indices,
                                            tf.shape(ref, out_type=tf.int64),
                                            update,
                                            default_value=0
                                            ),
                         dtype=tf.int64
                         )
  return ref * zero + update_value
//  [python]
def _build_update_func(self):
        @tfe.local_computation(player_name=self.player.name)
        def update_func(*grads):
            weights = self.model.weights
            grads = [t.cast(grad, t.float32) for grad in grads]
            with tf.name_scope("update_model"):
                for weight, grad in zip(weights, grads):
                    weight.assign(weight.read_value() - grad * self.LEARNING_RATE)

        self.update_func = update_func
//  [python]
def _build_update_func(self):
        @tfe.local_computation(player_name=self.player.name)
        def update_func(*grads):
            weights = self.model.weights
            grads = [tf.cast(grad, tf.float32) for grad in grads]
            with tf.name_scope("update_model"):
                for weight, grad in zip(weights, grads):
                    weight.assign(weight.read_value() - grad * self.LEARNING_RATE)

        self.update_func = update_func
//  [python]
def scale_bboxes(bbox, img_shape):
  """Scale bboxes to [0, 1). bbox format [ymin, xmin, ymax, xmax]
  Args:
    bbox: 2-D with shape '[num_bbox, 4]'
    img_shape: 1-D with shape '[4]'
  Return:
    sclaed_bboxes: scaled bboxes
  """
  img_h = t.cast(img_shape[0], dtype=t.float32)
  img_w = t.cast(img_shape[1], dtype=t.float32)
  shape = bbox.get_shape().as_list()
  _axis = 1 if len(shape) > 1 else 0
  [y_min, x_min, y_max, x_max] = tf.unstack(bbox, axis=_axis)
  y_1 = tf.truediv(y_min, img_h)
  x_1 = tf.truediv(x_min, img_w)
  y_2 = tf.truediv(y_max, img_h)
  x_2 = tf.truediv(x_max, img_w)
  return tf.stack([y_1, x_1, y_2, x_2], axis=_axis)
//  [python]
def scale_bboxes(bbox, img_shape):
  """Scale bboxes to [0, 1). bbox format [ymin, xmin, ymax, xmax]
  Args:
    bbox: 2-D with shape '[num_bbox, 4]'
    img_shape: 1-D with shape '[4]'
  Return:
    sclaed_bboxes: scaled bboxes
  """
  img_h = tf.cast(img_shape[0], dtype=tf.float32)
  img_w = tf.cast(img_shape[1], dtype=tf.float32)
  shape = bbox.get_shape().as_list()
  _axis = 1 if len(shape) > 1 else 0
  [y_min, x_min, y_max, x_max] = tf.unstack(bbox, axis=_axis)
  y_1 = tf.truediv(y_min, img_h)
  x_1 = tf.truediv(x_min, img_w)
  y_2 = tf.truediv(y_max, img_h)
  x_2 = tf.truediv(x_max, img_w)
  return tf.stack([y_1, x_1, y_2, x_2], axis=_axis)
//  [python]
def __next__(self):
        with tf.device('/cpu:0'):
            num = 0
            if self._batch_id == self._num_batches:
                self._batch_id = 0
                np.random.shuffle(self._image_ids)

            start = self._batch_size * (self._batch_id * self._num_threads + self._device_id)
            image_ids = self._image_ids[start : start + self._batch_size]

            images, bboxes, classes = self._input(image_ids)
            if self._is_training:
                self._color_twist(images)
                self._flip(images, bboxes)
                if self._use_mosaic:
                    images, bboxes, classes = self._mosaic(images, bboxes, classes)

            self._batch_id += 1

            lengths = [len(b) for b in bboxes]
            bboxes = tf.RaggedTensor.from_row_lengths(tf.concat(bboxes, axis=0), lengths)
            bboxes = bboxes.to_tensor(-1)
            bboxes = t.cast(bboxes, dtype=tf.float32)

            classes = tf.ragged.stack(classes)
            if self._batch_size > 1:
                classes = classes.to_tensor(-1)
            classes = t.cast(tf.expand_dims(classes, axis=-1), dtype=tf.float32)

            return images, tf.concat([bboxes, classes], axis=-1)
//  [python]
def __next__(self):
        with tf.device('/cpu:0'):
            num = 0
            if self._batch_id == self._num_batches:
                self._batch_id = 0
                np.random.shuffle(self._image_ids)

            start = self._batch_size * (self._batch_id * self._num_threads + self._device_id)
            image_ids = self._image_ids[start : start + self._batch_size]

            images, bboxes, classes = self._input(image_ids)
            if self._is_training:
                self._color_twist(images)
                self._flip(images, bboxes)
                if self._use_mosaic:
                    images, bboxes, classes = self._mosaic(images, bboxes, classes)

            self._batch_id += 1

            lengths = [len(b) for b in bboxes]
            bboxes = tf.RaggedTensor.from_row_lengths(tf.concat(bboxes, axis=0), lengths)
            bboxes = bboxes.to_tensor(-1)
            bboxes = tf.cast(bboxes, dtype=tf.float32)

            classes = tf.ragged.stack(classes)
            if self._batch_size > 1:
                classes = classes.to_tensor(-1)
            classes = tf.cast(tf.expand_dims(classes, axis=-1), dtype=tf.float32)

            return images, tf.concat([bboxes, classes], axis=-1)