jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def normalize_boxes(boxes, image_shape):
  """Converts boxes to the normalized coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    normalized_boxes: a tensor whose shape is the same as `boxes` representing
      the normalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError(
        'boxes.shape[-1] is {:d}, but must be 4.'.format(boxes.shape[-1]))

  with tf.name_scope('normalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.cast(image_shape, dtype=boxes.dtype)
      height = image_shape[..., 0:1]
      width = image_shape[..., 1:2]

    ymin = boxes[..., 0:1] / height
    xmin = boxes[..., 1:2] / width
    ymax = boxes[..., 2:3] / height
    xmax = boxes[..., 3:4] / width

    normalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return normalized_boxes
//  [python]
def __call__(self, class_outputs, class_targets, class_weights=None):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is categorical (or binary) cross entropy on all
    RoIs.
    Reference:
    https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py
    # pylint: disable=line-too-long

    Args:
      class_outputs: a float tensor representing the class prediction for each
        box with a shape of [batch_size, num_boxes, num_classes].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].
      class_weights: A float list containing the weight of each class.

    Returns:
      a scalar tensor representing total class loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      output_dtype = class_outputs.dtype
      num_classes = class_outputs.get_shape().as_list()[-1]
      class_weights = (
          class_weights if class_weights is not None else [1.0] * num_classes
      )
      if num_classes != len(class_weights):
        raise ValueError(
            'Length of class_weights should be {}'.format(num_classes)
        )

      class_weights = tf.constant(class_weights, dtype=output_dtype)

      class_targets_one_hot = tf.one_hot(
          tf.to_int32(class_targets),
          num_classes,
          dtype=class_outputs.dtype)
      if self._use_binary_cross_entropy:
        # (batch_size, num_boxes, num_classes)
        cross_entropy_loss = tf.nn.sigmoid_cross_entropy_with_logits(
            labels=class_targets_one_hot, logits=class_outputs)
        cross_entropy_loss *= class_weights
      else:
        # (batch_size, num_boxes)
        cross_entropy_loss = tf.nn.softmax_cross_entropy_with_logits(
            labels=class_targets_one_hot, logits=class_outputs)
        class_weight_mask = tf.einsum(
            '...y,y->...', class_targets_one_hot, class_weights
        )
        cross_entropy_loss *= class_weight_mask

      if self._top_k_percent < 1.0:
        return self.aggregate_loss_top_k(cross_entropy_loss)
      else:
        return tf.reduce_mean(cross_entropy_loss)
//  [python]
def __call__(self, class_outputs, class_targets, class_weights=None):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is categorical (or binary) cross entropy on all
    RoIs.
    Reference:
    https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py
    # pylint: disable=line-too-long

    Args:
      class_outputs: a float tensor representing the class prediction for each
        box with a shape of [batch_size, num_boxes, num_classes].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].
      class_weights: A float list containing the weight of each class.

    Returns:
      a scalar tensor representing total class loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      output_dtype = class_outputs.dtype
      num_classes = class_outputs.get_shape().as_list()[-1]
      class_weights = (
          class_weights if class_weights is not None else [1.0] * num_classes
      )
      if num_classes != len(class_weights):
        raise ValueError(
            'Length of class_weights should be {}'.format(num_classes)
        )

      class_weights = tf.constant(class_weights, dtype=output_dtype)

      class_targets_one_hot = tf.one_hot(
          tf.cast(class_targets, dtype=tf.int32),
          num_classes,
          dtype=class_outputs.dtype)
      if self._use_binary_cross_entropy:
        # (batch_size, num_boxes, num_classes)
        cross_entropy_loss = tf.nn.sigmoid_cross_entropy_with_logits(
            labels=class_targets_one_hot, logits=class_outputs)
        cross_entropy_loss *= class_weights
      else:
        # (batch_size, num_boxes)
        cross_entropy_loss = tf.nn.softmax_cross_entropy_with_logits(
            labels=class_targets_one_hot, logits=class_outputs)
        class_weight_mask = tf.einsum(
            '...y,y->...', class_targets_one_hot, class_weights
        )
        cross_entropy_loss *= class_weight_mask

      if self._top_k_percent < 1.0:
        return self.aggregate_loss_top_k(cross_entropy_loss)
      else:
        return tf.reduce_mean(cross_entropy_loss)
//  [python]
def clip_boxes(boxes, image_shape):
  """Clips boxes to image boundaries.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    clipped_boxes: a tensor whose shape is the same as `boxes` representing the
      clipped boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError(
        'boxes.shape[-1] is {:d}, but must be 4.'.format(boxes.shape[-1]))

  with tf.name_scope('clip_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
      max_length = [height, width, height, width]
    else:
      image_shape = tf.to_float(image_shape)
      height, width = tf.unstack(image_shape, axis=-1)
      max_length = tf.stack([height, width, height, width], axis=-1)

    clipped_boxes = tf.math.maximum(tf.math.minimum(boxes, max_length), 0.0)
    return clipped_boxes
//  [python]
def clip_boxes(boxes, image_shape):
  """Clips boxes to image boundaries.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    clipped_boxes: a tensor whose shape is the same as `boxes` representing the
      clipped boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError(
        'boxes.shape[-1] is {:d}, but must be 4.'.format(boxes.shape[-1]))

  with tf.name_scope('clip_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
      max_length = [height, width, height, width]
    else:
      image_shape = tf.cast(image_shape, dtype=boxes.dtype)
      height, width = tf.unstack(image_shape, axis=-1)
      max_length = tf.stack([height, width, height, width], axis=-1)

    clipped_boxes = tf.math.maximum(tf.math.minimum(boxes, max_length), 0.0)
    return clipped_boxes
//  [python]
def _mask(self, tensor, length):
    """Set padding elements of a batch of sequences to zero.

    Useful to then safely sum along the time dimension.

    Args:
      tensor: Tensor of sequences.
      length: Batch of sequence lengths.

    Returns:
      Masked sequences.
    """
    with tf.name_scope('mask'):
      range_ = tf.range(tensor.shape[1].value)
      mask = tf.to_float(range_[None, :] < length[:, None])
      masked = tensor * mask
      return tf.check_numerics(masked, 'masked')
//  [python]
def _mask(self, tensor, length):
    """Set padding elements of a batch of sequences to zero.

    Useful to then safely sum along the time dimension.

    Args:
      tensor: Tensor of sequences.
      length: Batch of sequence lengths.

    Returns:
      Masked sequences.
    """
    with tf.name_scope('mask'):
      range_ = tf.range(tensor.shape[1].value)
      mask = tf.cast(range_[None, :] < length[:, None], tf.float32)
      masked = tensor * mask
      return tf.check_numerics(masked, 'masked')
//  [python]
def encode_boxes(boxes, anchors, weights=None):
  """Encodes boxes to targets.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    encoded_boxes: a tensor whose shape is the same as `boxes` representing the
      encoded box targets.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError(
        'boxes.shape[-1] is {:d}, but must be 4.'.format(boxes.shape[-1]))

  with tf.name_scope('encode_boxes'):
    boxes = tf.to_float(boxes, dtype=anchors.dtype)
    ymin = boxes[..., 0:1]
    xmin = boxes[..., 1:2]
    ymax = boxes[..., 2:3]
    xmax = boxes[..., 3:4]
    box_h = ymax - ymin
    box_w = xmax - xmin
    box_yc = ymin + 0.5 * box_h
    box_xc = xmin + 0.5 * box_w

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]
    anchor_h = anchor_ymax - anchor_ymin
    anchor_w = anchor_xmax - anchor_xmin
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    # Avoid inf in log below.
    anchor_h += EPSILON
    anchor_w += EPSILON
    box_h += EPSILON
    box_w += EPSILON

    encoded_dy = (box_yc - anchor_yc) / anchor_h
    encoded_dx = (box_xc - anchor_xc) / anchor_w
    encoded_dh = tf.math.log(box_h / anchor_h)
    encoded_dw = tf.math.log(box_w / anchor_w)
    if weights:
      encoded_dy *= weights[0]
      encoded_dx *= weights[1]
      encoded_dh *= weights[2]
      encoded_dw *= weights[3]

    encoded_boxes = tf.concat(
        [encoded_dy, encoded_dx, encoded_dh, encoded_dw], axis=-1)
    return encoded_boxes
//  [python]
def encode_boxes(boxes, anchors, weights=None):
  """Encodes boxes to targets.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    encoded_boxes: a tensor whose shape is the same as `boxes` representing the
      encoded box targets.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError(
        'boxes.shape[-1] is {:d}, but must be 4.'.format(boxes.shape[-1]))

  with tf.name_scope('encode_boxes'):
    boxes = tf.cast(boxes, dtype=anchors.dtype)
    ymin = boxes[..., 0:1]
    xmin = boxes[..., 1:2]
    ymax = boxes[..., 2:3]
    xmax = boxes[..., 3:4]
    box_h = ymax - ymin
    box_w = xmax - xmin
    box_yc = ymin + 0.5 * box_h
    box_xc = xmin + 0.5 * box_w

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]
    anchor_h = anchor_ymax - anchor_ymin
    anchor_w = anchor_xmax - anchor_xmin
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    # Avoid inf in log below.
    anchor_h += EPSILON
    anchor_w += EPSILON
    box_h += EPSILON
    box_w += EPSILON

    encoded_dy = (box_yc - anchor_yc) / anchor_h
    encoded_dx = (box_xc - anchor_xc) / anchor_w
    encoded_dh = tf.math.log(box_h / anchor_h)
    encoded_dw = tf.math.log(box_w / anchor_w)
    if weights:
      encoded_dy *= weights[0]
      encoded_dx *= weights[1]
      encoded_dh *= weights[2]
      encoded_dw *= weights[3]

    encoded_boxes = tf.concat(
        [encoded_dy, encoded_dx, encoded_dh, encoded_dw], axis=-1)
    return encoded_boxes
//  [python]
def filter_boxes_by_scores(boxes, scores, min_score_threshold):
  """Filters and remove boxes whose scores are smaller than the threshold.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    scores: a tensor whose shape is the same as tf.shape(boxes)[:-1]
      representing the original scores of the boxes.
    min_score_threshold: a float representing the minimal box score threshold.
      Boxes whose score are smaller than it will be filtered out.

  Returns:
    filtered_boxes: a tensor whose shape is the same as `boxes` but with
      the position of the filtered boxes are filled with -1.
    filtered_scores: a tensor whose shape is the same as 'scores' but with
      the
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('filter_boxes_by_scores'):
    filtered_mask = tf.math.greater(scores, min_score_threshold)
    filtered_scores = tf.where(filtered_mask, scores, -tf.ones_like(scores))
    filtered_boxes = tf.to_float(
        tf.expand_dims(filtered_mask, axis=-1)) * boxes

    return filtered_boxes, filtered_scores
//  [python]
def filter_boxes_by_scores(boxes, scores, min_score_threshold):
  """Filters and remove boxes whose scores are smaller than the threshold.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    scores: a tensor whose shape is the same as tf.shape(boxes)[:-1]
      representing the original scores of the boxes.
    min_score_threshold: a float representing the minimal box score threshold.
      Boxes whose score are smaller than it will be filtered out.

  Returns:
    filtered_boxes: a tensor whose shape is the same as `boxes` but with
      the position of the filtered boxes are filled with -1.
    filtered_scores: a tensor whose shape is the same as 'scores' but with
      the
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('filter_boxes_by_scores'):
    filtered_mask = tf.math.greater(scores, min_score_threshold)
    filtered_scores = tf.where(filtered_mask, scores, -tf.ones_like(scores))
    filtered_boxes = tf.cast(
        tf.expand_dims(filtered_mask, axis=-1), dtype=boxes.dtype) * boxes

    return filtered_boxes, filtered_scores
//  [python]
def __call__(self, mask_outputs, mask_targets, select_class_targets):
    """Computes the mask loss of Mask-RCNN.

    This function implements the mask loss of Mask-RCNN. As the `mask_outputs`
    produces `num_classes` masks for each RoI, the reference model expands
    `mask_targets` to match the shape of `mask_outputs` and selects only the
    target that the RoI has a maximum overlap. (Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py)  # pylint: disable=line-too-long
    Instead, this implementation selects the `mask_outputs` by the
    `class_targets` so that it doesn't expand `mask_targets`. Note that the
    selection logic is done in the post-processing of mask_rcnn_fn in
    mask_rcnn_architecture.py.

    Args:
      mask_outputs: a float tensor representing the prediction for each mask,
        with a shape of
        [batch_size, num_masks, mask_height, mask_width].
      mask_targets: a float tensor representing the binary mask of ground truth
        labels for each mask with a shape of
        [batch_size, num_masks, mask_height, mask_width].
      select_class_targets: a tensor with a shape of [batch_size, num_masks],
        representing the foreground mask targets.

    Returns:
      mask_loss: a float tensor representing total mask loss.
    """
    with tf.name_scope('mask_rcnn_loss'):
      _, _, mask_height, mask_width = mask_outputs.get_shape().as_list()

      weights = tf.tile(
          tf.greater(select_class_targets, 0)[:, :, tf.newaxis, tf.newaxis],
          [1, 1, mask_height, mask_width],
      )
      weights = tf.to_float(weights)

      mask_targets = tf.expand_dims(mask_targets, axis=-1)
      mask_outputs = tf.expand_dims(mask_outputs, axis=-1)
      mask_loss = self._binary_crossentropy(mask_targets, mask_outputs,
                                            sample_weight=weights)
      # For each replica, get the sum of non-zero weights. Then get the mean of
      # sums from all replicas. Note there is an extra division by
      # `num_replicas` in train_step(). So it is equivalent to normalizing the
      # mask loss by the global sum of non-zero weights.
      replica_context = tf.distribute.get_replica_context()
      weights = tf.reduce_sum(weights)
      weights_mean = replica_context.all_reduce(
          tf.distribute.ReduceOp.MEAN, weights
      )
      # The loss is normalized by the number of 1's in weights and
      # + 0.01 is used to avoid division by zero.
      return mask_loss / (weights_mean + 0.01)
//  [python]
def __call__(self, mask_outputs, mask_targets, select_class_targets):
    """Computes the mask loss of Mask-RCNN.

    This function implements the mask loss of Mask-RCNN. As the `mask_outputs`
    produces `num_classes` masks for each RoI, the reference model expands
    `mask_targets` to match the shape of `mask_outputs` and selects only the
    target that the RoI has a maximum overlap. (Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py)  # pylint: disable=line-too-long
    Instead, this implementation selects the `mask_outputs` by the
    `class_targets` so that it doesn't expand `mask_targets`. Note that the
    selection logic is done in the post-processing of mask_rcnn_fn in
    mask_rcnn_architecture.py.

    Args:
      mask_outputs: a float tensor representing the prediction for each mask,
        with a shape of
        [batch_size, num_masks, mask_height, mask_width].
      mask_targets: a float tensor representing the binary mask of ground truth
        labels for each mask with a shape of
        [batch_size, num_masks, mask_height, mask_width].
      select_class_targets: a tensor with a shape of [batch_size, num_masks],
        representing the foreground mask targets.

    Returns:
      mask_loss: a float tensor representing total mask loss.
    """
    with tf.name_scope('mask_rcnn_loss'):
      _, _, mask_height, mask_width = mask_outputs.get_shape().as_list()

      weights = tf.tile(
          tf.greater(select_class_targets, 0)[:, :, tf.newaxis, tf.newaxis],
          [1, 1, mask_height, mask_width],
      )
      weights = tf.cast(weights, dtype=mask_outputs.dtype)

      mask_targets = tf.expand_dims(mask_targets, axis=-1)
      mask_outputs = tf.expand_dims(mask_outputs, axis=-1)
      mask_loss = self._binary_crossentropy(mask_targets, mask_outputs,
                                            sample_weight=weights)
      # For each replica, get the sum of non-zero weights. Then get the mean of
      # sums from all replicas. Note there is an extra division by
      # `num_replicas` in train_step(). So it is equivalent to normalizing the
      # mask loss by the global sum of non-zero weights.
      replica_context = tf.distribute.get_replica_context()
      weights = tf.reduce_sum(weights)
      weights_mean = replica_context.all_reduce(
          tf.distribute.ReduceOp.MEAN, weights
      )
      # The loss is normalized by the number of 1's in weights and
      # + 0.01 is used to avoid division by zero.
      return mask_loss / (weights_mean + 0.01)
//  [python]
def build_dataset(self, filenames, labels, is_training):
    """Build input dataset."""
    filenames = tf.constant(filenames)
    labels = tf.constant(labels)

    dataset = tf.compat.v1.data.Dataset.from_tensor_slices((filenames, labels))

    def _parse_function(filename, label):
      image_string = tf.io.read_file(filename)
      image_decoded = preprocessing.preprocess_image(
          image_string, is_training, self.image_size)
      image = tf.compat.v1.cast(image_decoded, tf.float32)
      return image, label

    dataset = dataset.map(_parse_function)
    dataset = dataset.batch(self.batch_size)

    iterator = dataset.make_one_shot_iterator()
    #iterator = iter(dataset)
    images, labels = iterator.get_next()
    return images, labels
//  [python]
def build_dataset(self, filenames, labels, is_training):
    """Build input dataset."""
    filenames = tf.constant(filenames)
    labels = tf.constant(labels)

    dataset = tf.compat.v1.data.Dataset.from_tensor_slices((filenames, labels))

    def _parse_function(filename, label):
      image_string = tf.io.read_file(filename)
      image_decoded = preprocessing.preprocess_image(
          image_string, is_training, self.image_size)
      image = tf.cast(image_decoded, tf.float32)
      return image, label

    dataset = dataset.map(_parse_function)
    dataset = dataset.batch(self.batch_size)

    iterator = dataset.make_one_shot_iterator()
    #iterator = iter(dataset)
    images, labels = iterator.get_next()
    return images, labels
//  [python]
def _nearest_indices_and_weights(coordinate):
    coordinate = (
        coordinate if coordinate.dtype.is_integer else tf.round(coordinate)
    )
    index = t.cast(coordinate, tf.int32)
    weight = tf.constant(1, coordinate.dtype)
    return [(index, weight)]
//  [python]
def _nearest_indices_and_weights(coordinate):
    coordinate = (
        coordinate if coordinate.dtype.is_integer else tf.round(coordinate)
    )
    index = tf.cast(coordinate, tf.int32)
    weight = tf.constant(1, coordinate.dtype)
    return [(index, weight)]
//  [python]
def call(self, inputs: tf.Tensor):
    """Calls `text.BertTokenizer` on inputs.

    Args:
      inputs: A string Tensor of shape [batch_size].

    Returns:
      One or three of `RaggedTensors` if `tokenize_with_offsets` is False or
      True, respectively. These are
      tokens: A `RaggedTensor` of shape [batch_size, (words), (pieces_per_word)]
        and type int32. tokens[i,j,k] contains the k-th wordpiece of the
        j-th word in the i-th input.
      start_offsets, limit_offsets: If `tokenize_with_offsets` is True,
        RaggedTensors of type int64 with the same indices as tokens.
        Element [i,j,k] contains the byte offset at the start, or past the
        end, resp., for the k-th wordpiece of the j-th word in the i-th input.
    """
    # Prepare to reshape the result to work around broken shape inference.
    batch_size = tf.shape(inputs)[0]
    def _reshape(rt):
      values = rt.values
      row_splits = rt.row_splits
      row_splits = tf.reshape(row_splits, [batch_size + 1])
      return tf.RaggedTensor.from_row_splits(values, row_splits)

    # Call the tokenizer.
    if self.tokenize_with_offsets:
      tokens, start_offsets, limit_offsets = (
          self._bert_tokenizer.tokenize_with_offsets(inputs))
      tokens = t.cast(tokens, dtype=tf.int32)
      return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
    else:
      tokens = self._bert_tokenizer.tokenize(inputs)
      tokens = t.cast(tokens, dtype=tf.int32)
      return _reshape(tokens)
//  [python]
def call(self, inputs: tf.Tensor):
    """Calls `text.BertTokenizer` on inputs.

    Args:
      inputs: A string Tensor of shape [batch_size].

    Returns:
      One or three of `RaggedTensors` if `tokenize_with_offsets` is False or
      True, respectively. These are
      tokens: A `RaggedTensor` of shape [batch_size, (words), (pieces_per_word)]
        and type int32. tokens[i,j,k] contains the k-th wordpiece of the
        j-th word in the i-th input.
      start_offsets, limit_offsets: If `tokenize_with_offsets` is True,
        RaggedTensors of type int64 with the same indices as tokens.
        Element [i,j,k] contains the byte offset at the start, or past the
        end, resp., for the k-th wordpiece of the j-th word in the i-th input.
    """
    # Prepare to reshape the result to work around broken shape inference.
    batch_size = tf.shape(inputs)[0]
    def _reshape(rt):
      values = rt.values
      row_splits = rt.row_splits
      row_splits = tf.reshape(row_splits, [batch_size + 1])
      return tf.RaggedTensor.from_row_splits(values, row_splits)

    # Call the tokenizer.
    if self.tokenize_with_offsets:
      tokens, start_offsets, limit_offsets = (
          self._bert_tokenizer.tokenize_with_offsets(inputs))
      tokens = tf.cast(tokens, dtype=tf.int32)
      return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
    else:
      tokens = self._bert_tokenizer.tokenize(inputs)
      tokens = tf.cast(tokens, dtype=tf.int32)
      return _reshape(tokens)
//  [python]
def result(self) -> tf.Tensor:
    """Computes IoU for each class."""
    return tf.cast(
        tf.math.divide_no_nan(self.intersection_per_class,
                              self.union_per_class), self.dtype)
//  [python]
def decode_boxes(encoded_boxes, anchors, weights=None):
  """Decodes boxes.

  Args:
    encoded_boxes: a tensor whose last dimension is 4 representing the
      coordinates of encoded boxes in dy, dx, dh, dw in order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    decoded_boxes: a tensor whose shape is the same as `boxes` representing the
      decoded box targets.
  """
  if encoded_boxes.shape[-1] != 4:
    raise ValueError(
        'encoded_boxes.shape[-1] is {:d}, but must be 4.'
        .format(encoded_boxes.shape[-1]))

  with tf.name_scope('decode_boxes'):
    encoded_boxes = tf.to_float(encoded_boxes)
    dy, dx, dh, dw = tf.split(encoded_boxes, 4, -1)
    if weights:
      dy /= weights[0]
      dx /= weights[1]
      dh /= weights[2]
      dw /= weights[3]
    dh = tf.math.minimum(dh, BBOX_XFORM_CLIP)
    dw = tf.math.minimum(dw, BBOX_XFORM_CLIP)

    anchor_ymin, anchor_xmin, anchor_ymax, anchor_xmax = tf.split(
        anchors, 4, -1)
    anchor_h = anchor_ymax - anchor_ymin
    anchor_w = anchor_xmax - anchor_xmin
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    decoded_boxes_yc = dy * anchor_h + anchor_yc
    decoded_boxes_xc = dx * anchor_w + anchor_xc
    decoded_boxes_h = tf.math.exp(dh) * anchor_h
    decoded_boxes_w = tf.math.exp(dw) * anchor_w

    decoded_boxes_ymin = decoded_boxes_yc - 0.5 * decoded_boxes_h
    decoded_boxes_xmin = decoded_boxes_xc - 0.5 * decoded_boxes_w
    decoded_boxes_ymax = decoded_boxes_ymin + decoded_boxes_h
    decoded_boxes_xmax = decoded_boxes_xmin + decoded_boxes_w

    decoded_boxes = tf.concat(
        [decoded_boxes_ymin, decoded_boxes_xmin,
         decoded_boxes_ymax, decoded_boxes_xmax],
        axis=-1)
    return decoded_boxes
//  [python]
def decode_boxes(encoded_boxes, anchors, weights=None):
  """Decodes boxes.

  Args:
    encoded_boxes: a tensor whose last dimension is 4 representing the
      coordinates of encoded boxes in dy, dx, dh, dw in order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    decoded_boxes: a tensor whose shape is the same as `boxes` representing the
      decoded box targets.
  """
  if encoded_boxes.shape[-1] != 4:
    raise ValueError(
        'encoded_boxes.shape[-1] is {:d}, but must be 4.'
        .format(encoded_boxes.shape[-1]))

  with tf.name_scope('decode_boxes'):
    encoded_boxes = tf.cast(encoded_boxes, dtype=anchors.dtype)
    dy, dx, dh, dw = tf.split(encoded_boxes, 4, -1)
    if weights:
      dy /= weights[0]
      dx /= weights[1]
      dh /= weights[2]
      dw /= weights[3]
    dh = tf.math.minimum(dh, BBOX_XFORM_CLIP)
    dw = tf.math.minimum(dw, BBOX_XFORM_CLIP)

    anchor_ymin, anchor_xmin, anchor_ymax, anchor_xmax = tf.split(
        anchors, 4, -1)
    anchor_h = anchor_ymax - anchor_ymin
    anchor_w = anchor_xmax - anchor_xmin
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    decoded_boxes_yc = dy * anchor_h + anchor_yc
    decoded_boxes_xc = dx * anchor_w + anchor_xc
    decoded_boxes_h = tf.math.exp(dh) * anchor_h
    decoded_boxes_w = tf.math.exp(dw) * anchor_w

    decoded_boxes_ymin = decoded_boxes_yc - 0.5 * decoded_boxes_h
    decoded_boxes_xmin = decoded_boxes_xc - 0.5 * decoded_boxes_w
    decoded_boxes_ymax = decoded_boxes_ymin + decoded_boxes_h
    decoded_boxes_xmax = decoded_boxes_xmin + decoded_boxes_w

    decoded_boxes = tf.concat(
        [decoded_boxes_ymin, decoded_boxes_xmin,
         decoded_boxes_ymax, decoded_boxes_xmax],
        axis=-1)
    return decoded_boxes
//  [python]
def rgb_to_grayscale(image, data_format="channels_last"):
    if data_format == "channels_first":
        if len(image.shape) == 4:
            image = tf.transpose(image, (0, 2, 3, 1))
        elif len(image.shape) == 3:
            image = tf.transpose(image, (1, 2, 0))
        else:
            raise ValueError(
                "Invalid input rank: expected rank 3 (single image) "
                "or rank 4 (batch of images). Received input with shape: "
                f"image.shape={image.shape}"
            )
    grayscale_image = tf.image.rgb_to_grayscale(image)
    if data_format == "channels_first":
        if len(image.shape) == 4:
            grayscale_image = tf.transpose(grayscale_image, (0, 3, 1, 2))
        elif len(image.shape) == 3:
            grayscale_image = tf.transpose(grayscale_image, (2, 0, 1))
    return tf.to_float(grayscale_image)
//  [python]
def rgb_to_grayscale(image, data_format="channels_last"):
    if data_format == "channels_first":
        if len(image.shape) == 4:
            image = tf.transpose(image, (0, 2, 3, 1))
        elif len(image.shape) == 3:
            image = tf.transpose(image, (1, 2, 0))
        else:
            raise ValueError(
                "Invalid input rank: expected rank 3 (single image) "
                "or rank 4 (batch of images). Received input with shape: "
                f"image.shape={image.shape}"
            )
    grayscale_image = tf.image.rgb_to_grayscale(image)
    if data_format == "channels_first":
        if len(image.shape) == 4:
            grayscale_image = tf.transpose(grayscale_image, (0, 3, 1, 2))
        elif len(image.shape) == 3:
            grayscale_image = tf.transpose(grayscale_image, (2, 0, 1))
    return tf.cast(grayscale_image, image.dtype)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                            np.arange(d_model)[np.newaxis, :],
                            d_model)

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encoding = angle_rads[np.newaxis, ...]

    return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                            np.arange(d_model)[np.newaxis, :],
                            d_model)

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encoding = angle_rads[np.newaxis, ...]

    return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def _policy_loss(self, mean, logstd, old_mean, old_logstd, action, advantage, length):
    """Compute the policy loss composed of multiple components.

    1. The policy gradient loss is importance sampled from the data-collecting
       policy at the beginning of training.
    2. The second term is a KL penalty between the policy at the beginning of
       training and the current policy.
    3. Additionally, if this KL already changed more than twice the target
       amount, we activate a strong penalty discouraging further divergence.

    Args:
      mean: Sequences of action means of the current policy.
      logstd: Sequences of action log stddevs of the current policy.
      old_mean: Sequences of action means of the behavioral policy.
      old_logstd: Sequences of action log stddevs of the behavioral policy.
      action: Sequences of actions.
      advantage: Sequences of advantages.
      length: Batch of sequence lengths.

    Returns:
      Tuple of loss tensor and summary tensor.
    """
    with tf.name_scope('policy_loss'):
      entropy = utility.diag_normal_entropy(mean, logstd)
      kl = tf.reduce_mean(
          self._mask(utility.diag_normal_kl(old_mean, old_logstd, mean, logstd), length), 1)
      policy_gradient = tf.exp(
          utility.diag_normal_logpdf(mean, logstd, action) -
          utility.diag_normal_logpdf(old_mean, old_logstd, action))
      surrogate_loss = -tf.reduce_mean(
          self._mask(policy_gradient * tf.stop_gradient(advantage), length), 1)
      kl_penalty = self._penalty * kl
      cutoff_threshold = self._config.kl_target * self._config.kl_cutoff_factor
      cutoff_count = tf.reduce_sum(tf.to_int32(tf.greater(kl, cutoff_threshold)))
      with tf.control_dependencies(
          [tf.cond(cutoff_count > 0, lambda: tf.Print(0, [cutoff_count], 'kl cutoff! '), int)]):
        kl_cutoff = (self._config.kl_cutoff_coef * tf.to_float(tf.greater(kl, cutoff_threshold)) *
                     (kl - cutoff_threshold)**2)
      policy_loss = surrogate_loss + kl_penalty + kl_cutoff
      summary = tf.summary.merge([
          tf.summary.histogram('entropy', entropy),
          tf.summary.histogram('kl', kl),
          tf.summary.histogram('surrogate_loss', surrogate_loss),
          tf.summary.histogram('kl_penalty', kl_penalty),
          tf.summary.histogram('kl_cutoff', kl_cutoff),
          tf.summary.histogram('kl_penalty_combined', kl_penalty + kl_cutoff),
          tf.summary.histogram('policy_loss', policy_loss),
          tf.summary.scalar('avg_surr_loss', tf.reduce_mean(surrogate_loss)),
          tf.summary.scalar('avg_kl_penalty', tf.reduce_mean(kl_penalty)),
          tf.summary.scalar('avg_policy_loss', tf.reduce_mean(policy_loss))
      ])
      policy_loss = tf.reduce_mean(policy_loss, 0)
      return tf.check_numerics(policy_loss, 'policy_loss'), summary
//  [python]
def _policy_loss(self, mean, logstd, old_mean, old_logstd, action, advantage, length):
    """Compute the policy loss composed of multiple components.

    1. The policy gradient loss is importance sampled from the data-collecting
       policy at the beginning of training.
    2. The second term is a KL penalty between the policy at the beginning of
       training and the current policy.
    3. Additionally, if this KL already changed more than twice the target
       amount, we activate a strong penalty discouraging further divergence.

    Args:
      mean: Sequences of action means of the current policy.
      logstd: Sequences of action log stddevs of the current policy.
      old_mean: Sequences of action means of the behavioral policy.
      old_logstd: Sequences of action log stddevs of the behavioral policy.
      action: Sequences of actions.
      advantage: Sequences of advantages.
      length: Batch of sequence lengths.

    Returns:
      Tuple of loss tensor and summary tensor.
    """
    with tf.name_scope('policy_loss'):
      entropy = utility.diag_normal_entropy(mean, logstd)
      kl = tf.reduce_mean(
          self._mask(utility.diag_normal_kl(old_mean, old_logstd, mean, logstd), length), 1)
      policy_gradient = tf.exp(
          utility.diag_normal_logpdf(mean, logstd, action) -
          utility.diag_normal_logpdf(old_mean, old_logstd, action))
      surrogate_loss = -tf.reduce_mean(
          self._mask(policy_gradient * tf.stop_gradient(advantage), length), 1)
      kl_penalty = self._penalty * kl
      cutoff_threshold = self._config.kl_target * self._config.kl_cutoff_factor
      cutoff_count = tf.reduce_sum(tf.cast(kl > cutoff_threshold, tf.int32))
      with tf.control_dependencies(
          [tf.cond(cutoff_count > 0, lambda: tf.Print(0, [cutoff_count], 'kl cutoff! '), int)]):
        kl_cutoff = (self._config.kl_cutoff_coef * tf.cast(kl > cutoff_threshold, tf.float32) *
                     (kl - cutoff_threshold)**2)
      policy_loss = surrogate_loss + kl_penalty + kl_cutoff
      summary = tf.summary.merge([
          tf.summary.histogram('entropy', entropy),
          tf.summary.histogram('kl', kl),
          tf.summary.histogram('surrogate_loss', surrogate_loss),
          tf.summary.histogram('kl_penalty', kl_penalty),
          tf.summary.histogram('kl_cutoff', kl_cutoff),
          tf.summary.histogram('kl_penalty_combined', kl_penalty + kl_cutoff),
          tf.summary.histogram('policy_loss', policy_loss),
          tf.summary.scalar('avg_surr_loss', tf.reduce_mean(surrogate_loss)),
          tf.summary.scalar('avg_kl_penalty', tf.reduce_mean(kl_penalty)),
          tf.summary.scalar('avg_policy_loss', tf.reduce_mean(policy_loss))
      ])
      policy_loss = tf.reduce_mean(policy_loss, 0)
      return tf.check_numerics(policy_loss, 'policy_loss'), summary
//  [python]
def call(self, x, training, mask):
        # all inp image sequences are always 49, so mask not needed
        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        # input size should be batch_size, 49, 2048)
        x = self.fc(x)  # output dims should be (batch_size, 49, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))  # scaled dot product attention
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](
                x, training, mask)  # mask shouldnt be needed

        return x
//  [python]
def update_state(self, y_true: tf.Tensor, y_pred: tf.Tensor):
    """Updates metric state by accumulating the variables.

    Args:
      y_true: The ground truth values.
      y_pred: The predicted values.
    """

    if self.sparse_y_true:
      # Shape: (..., num_classes, ...)
      y_true = tf.one_hot(
          tf.to_int32(y_true),
          self.num_classes,
          axis=self.axis,
          on_value=True,
          off_value=False,
      )
    if self.sparse_y_pred:
      # Shape: (..., num_classes, ...)
      y_pred = tf.one_hot(
          tf.to_int32(y_pred),
          self.num_classes,
          axis=self.axis,
          on_value=True,
          off_value=False,
      )

    one_hot_axis = self.axis if self.axis >= 0 else (
        len(y_true.get_shape().as_list()) + self.axis)
    # Reduce sum the leading dimensions.
    # Shape: (num_classes, ...)
    current_intersection = tf.math.count_nonzero(
        y_pred & y_true, axis=np.arange(one_hot_axis), dtype=tf.float32
    )
    # Shape: (num_classes, ...)
    current_union = tf.math.count_nonzero(
        y_pred | y_true, axis=np.arange(one_hot_axis), dtype=tf.float32
    )

    self.intersection_per_class.assign_add(
        tf.to_float(current_intersection))
    self.union_per_class.assign_add(
        tf.to_float(current_union))
//  [python]
def update_state(self, y_true: tf.Tensor, y_pred: tf.Tensor):
    """Updates metric state by accumulating the variables.

    Args:
      y_true: The ground truth values.
      y_pred: The predicted values.
    """

    if self.sparse_y_true:
      # Shape: (..., num_classes, ...)
      y_true = tf.one_hot(
          tf.cast(y_true, dtype=tf.int32),
          self.num_classes,
          axis=self.axis,
          on_value=True,
          off_value=False,
      )
    if self.sparse_y_pred:
      # Shape: (..., num_classes, ...)
      y_pred = tf.one_hot(
          tf.cast(y_pred, dtype=tf.int32),
          self.num_classes,
          axis=self.axis,
          on_value=True,
          off_value=False,
      )

    one_hot_axis = self.axis if self.axis >= 0 else (
        len(y_true.get_shape().as_list()) + self.axis)
    # Reduce sum the leading dimensions.
    # Shape: (num_classes, ...)
    current_intersection = tf.math.count_nonzero(
        y_pred & y_true, axis=np.arange(one_hot_axis), dtype=tf.float32
    )
    # Shape: (num_classes, ...)
    current_union = tf.math.count_nonzero(
        y_pred | y_true, axis=np.arange(one_hot_axis), dtype=tf.float32
    )

    self.intersection_per_class.assign_add(
        tf.cast(current_intersection, self.intersection_per_class.dtype))
    self.union_per_class.assign_add(
        tf.cast(current_union, self.union_per_class.dtype))
//  [python]
def call(self, x, training, mask):
        # we need to convert 224x224x3 images in 49x2048
        x = self.features_extract(x)
        x = tf.reshape(x,
                    (x.shape[0], -1, x.shape[3])) # 49x2048 now
        
        #all inp image sequences are always 49, so mask not needed
        seq_len = tf.shape(x)[1] 
        
        #x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        
        # adding embedding and position encoding.
        # input size should be batch_size, 49, 2048)
        x = self.fc(x)  # output dims should be (batch_size, 49, d_model)
        x *= tf.to_float(self.d_model)
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)
        
        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)  # mask shouldnt be needed
        
        return x
//  [python]
def call(self, x, training, mask):
        # we need to convert 224x224x3 images in 49x2048
        x = self.features_extract(x)
        x = tf.reshape(x,
                    (x.shape[0], -1, x.shape[3])) # 49x2048 now
        
        #all inp image sequences are always 49, so mask not needed
        seq_len = tf.shape(x)[1] 
        
        #x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        
        # adding embedding and position encoding.
        # input size should be batch_size, 49, 2048)
        x = self.fc(x)  # output dims should be (batch_size, 49, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)
        
        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)  # mask shouldnt be needed
        
        return x
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead)
    but it must be broadcastable for addition.

    Args:
      q: query shape == (..., seq_len_q, depth)
      k: key shape == (..., seq_len_k, depth)
      v: value shape == (..., seq_len_v, depth_v)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    """

    # (..., seq_len_q, seq_len_k)
    matmul_qk = tf.matmul(q, k, transpose_b=True)

    # scale matmul_qk
    dk = tf.to_float(tf.shape(k)[-1])
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(
        scaled_attention_logits,
        axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead)
    but it must be broadcastable for addition.

    Args:
      q: query shape == (..., seq_len_q, depth)
      k: key shape == (..., seq_len_k, depth)
      v: value shape == (..., seq_len_v, depth_v)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    """

    # (..., seq_len_q, seq_len_k)
    matmul_qk = tf.matmul(q, k, transpose_b=True)

    # scale matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(
        scaled_attention_logits,
        axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights
//  [python]
def call(self, x, enc_output, training,
             look_ahead_mask, padding_mask):

        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                                   look_ahead_mask, 
                                                   padding_mask)

        attention_weights['decoder_layer{}_block1'.format(i + 1)] = block1
        attention_weights['decoder_layer{}_block2'.format(i + 1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights
//  [python]
def from_4d(image: tf.Tensor, ndims: tf.Tensor) -> tf.Tensor:
  """Converts a 4D image back to `ndims` rank."""
  shape = tf.shape(image)
  begin = tf.to_int32(tf.less_equal(ndims, 3))
  end = 4 - tf.to_int32(tf.equal(ndims, 2))
  new_shape = shape[begin:end]
  return tf.reshape(image, new_shape)
//  [python]
def from_4d(image: tf.Tensor, ndims: tf.Tensor) -> tf.Tensor:
  """Converts a 4D image back to `ndims` rank."""
  shape = tf.shape(image)
  begin = tf.cast(tf.less_equal(ndims, 3), dtype=tf.int32)
  end = 4 - tf.cast(tf.equal(ndims, 2), dtype=tf.int32)
  new_shape = shape[begin:end]
  return tf.reshape(image, new_shape)