jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(
        self,
        inputs: Union[tf.Tensor, TabularData],
        training: bool = True,
        testing: bool = False,
        **kwargs,
    ) -> Union[tf.Tensor, TabularData]:
        """Based on the user/query embedding (inputs[self.query_name]), uses dot product to score
            the positive item (inputs["item"]).
            For the sampled-softmax mode, logits are computed by multiplying the query vector
            and the item embeddings matrix (self.context.get_embedding(self.item_domain))

        Parameters
        ----------
        inputs : Union[tf.Tensor, TabularData]
            Dict with the query and item embeddings (e.g. `{"query": <emb>}, "item": <emb>}`),
            where embeddings are 2D tensors (batch size, embedding size)
        training : bool, optional
            Flag that indicates whether in training mode, by default True

        Returns
        -------
        tf.Tensor
            2D Tensor with the scores for the positive items,
            If `training=True`, return the original inputs
        """
        if self.cache_query:
            # enabled only during top-k evaluation

            query = inputs[self.query_name]
            context_query_size = tf.shape(self.context["query"])[0]
            # pad with zeros to match shape of initial query variable
            padding_size = context_query_size - tf.shape(query)[0]
            if padding_size > 0:
                query = tf.pad(query, [[0, padding_size], [0, 0]])

            query = query[:context_query_size]

            self.context["query"].assign(t.cast(query, tf.float32))

        if training or testing:
            return inputs

        if self.sampled_softmax_mode:
            return self._get_logits_for_sampled_softmax(inputs)

        self._check_input_from_two_tower(inputs)
        positive_scores = tf.reduce_sum(
            tf.multiply(inputs[self.query_name], inputs[self.item_name]), keepdims=True, axis=-1
        )
        return positive_scores
//  [python]
def call(
        self,
        inputs: Union[tf.Tensor, TabularData],
        training: bool = True,
        testing: bool = False,
        **kwargs,
    ) -> Union[tf.Tensor, TabularData]:
        """Based on the user/query embedding (inputs[self.query_name]), uses dot product to score
            the positive item (inputs["item"]).
            For the sampled-softmax mode, logits are computed by multiplying the query vector
            and the item embeddings matrix (self.context.get_embedding(self.item_domain))

        Parameters
        ----------
        inputs : Union[tf.Tensor, TabularData]
            Dict with the query and item embeddings (e.g. `{"query": <emb>}, "item": <emb>}`),
            where embeddings are 2D tensors (batch size, embedding size)
        training : bool, optional
            Flag that indicates whether in training mode, by default True

        Returns
        -------
        tf.Tensor
            2D Tensor with the scores for the positive items,
            If `training=True`, return the original inputs
        """
        if self.cache_query:
            # enabled only during top-k evaluation

            query = inputs[self.query_name]
            context_query_size = tf.shape(self.context["query"])[0]
            # pad with zeros to match shape of initial query variable
            padding_size = context_query_size - tf.shape(query)[0]
            if padding_size > 0:
                query = tf.pad(query, [[0, padding_size], [0, 0]])

            query = query[:context_query_size]

            self.context["query"].assign(tf.cast(query, tf.float32))

        if training or testing:
            return inputs

        if self.sampled_softmax_mode:
            return self._get_logits_for_sampled_softmax(inputs)

        self._check_input_from_two_tower(inputs)
        positive_scores = tf.reduce_sum(
            tf.multiply(inputs[self.query_name], inputs[self.item_name]), keepdims=True, axis=-1
        )
        return positive_scores
//  [python]
def test_pick_fourier_transform(self):
    # Ensure we don't hit an edge case which exceeds the fixed numerical error.
    tf.random.set_seed(1)
    np.random.seed(1)

    batch_size = 3
    max_seq_length = 4
    hidden_dim = 8

    fft = mixing._pick_fourier_transform(
        use_fft=True, max_seq_length=max_seq_length, hidden_dim=hidden_dim)
    dft_matmul = mixing._pick_fourier_transform(
        use_fft=False, max_seq_length=max_seq_length, hidden_dim=hidden_dim)

    inputs = tf.random.uniform([batch_size, max_seq_length, hidden_dim])
    inputs = tf.cast(inputs, tf.complex64)

    np.testing.assert_allclose(
        fft(inputs), dft_matmul(inputs), rtol=1e-6, atol=1e-6)
//  [python]
def normalize_states(self, states):
    dtype = tf.convert_to_tensor(states).dtype
    return ((states - self.state_mean) /
            tf.maximum(tf.cast(self.eps, dtype), self.state_std))
//  [python]
def unnormalize_states(self, states):
    dtype = tf.convert_to_tensor(states).dtype
    return (states * tf.maximum(tf.cast(self.eps, dtype), self.state_std)
            + self.state_mean)
//  [python]
def _get_factorization(self, inputs: tf.Tensor, input_mask: tf.Tensor):
    """Samples a permutation of the factorization order.

    Args:
      inputs: the input tokens.
      input_mask: the `bool` Tensor of the same shape as `inputs`. If `True`,
        then this means select for partial prediction.

    Returns:
      perm_mask: An `int32` Tensor of shape [seq_length, seq_length] consisting
        of 0s and 1s. If perm_mask[i][j] == 0, then this means that the i-th
        token (in original order) cannot attend to the jth attention token.
      target_mask: An `int32` Tensor of shape [seq_len] consisting of 0s and 1s.
        If target_mask[i] == 1, then the i-th token needs to be predicted and
        the mask will be used as input. This token will be included in the loss.
        If target_mask[i] == 0, then the token (or [SEP], [CLS]) will be used as
        input. This token will not be included in the loss.
      tokens: int32 Tensor of shape [seq_length].
      masked_tokens: int32 Tensor of shape [seq_length].
    """
    factorization_length = tf.shape(inputs)[0]
    # Generate permutation indices
    index = tf.range(factorization_length, dtype=tf.int32)
    index = tf.transpose(tf.reshape(index, [-1, self._permutation_size]))
    index = tf.random.shuffle(index)
    index = tf.reshape(tf.transpose(index), [-1])

    input_mask = t.cast(input_mask, tf.bool)

    # non-functional tokens
    non_func_tokens = tf.logical_not(
        tf.logical_or(
            tf.equal(inputs, self._sep_id), tf.equal(inputs, self._cls_id)))
    masked_tokens = tf.logical_and(input_mask, non_func_tokens)
    non_masked_or_func_tokens = tf.logical_not(masked_tokens)

    smallest_index = -2 * tf.ones([factorization_length], dtype=tf.int32)

    # Similar to BERT, randomly leak some masked tokens
    if self._leak_ratio > 0:
      leak_tokens = tf.logical_and(
          masked_tokens,
          tf.random.uniform([factorization_length], maxval=1.0) <
          self._leak_ratio)
      can_attend_self = tf.logical_or(non_masked_or_func_tokens, leak_tokens)
    else:
      can_attend_self = non_masked_or_func_tokens
    to_index = tf.where(can_attend_self, smallest_index, index)
    from_index = tf.where(can_attend_self, to_index + 1, to_index)

    # For masked tokens, can attend if i > j
    # For context tokens, always can attend each other
    can_attend = from_index[:, None] > to_index[None, :]

    perm_mask = t.cast(can_attend, tf.int32)

    # Only masked tokens are included in the loss
    target_mask = t.cast(masked_tokens, tf.int32)

    return perm_mask, target_mask, inputs, masked_tokens
//  [python]
def _get_factorization(self, inputs: tf.Tensor, input_mask: tf.Tensor):
    """Samples a permutation of the factorization order.

    Args:
      inputs: the input tokens.
      input_mask: the `bool` Tensor of the same shape as `inputs`. If `True`,
        then this means select for partial prediction.

    Returns:
      perm_mask: An `int32` Tensor of shape [seq_length, seq_length] consisting
        of 0s and 1s. If perm_mask[i][j] == 0, then this means that the i-th
        token (in original order) cannot attend to the jth attention token.
      target_mask: An `int32` Tensor of shape [seq_len] consisting of 0s and 1s.
        If target_mask[i] == 1, then the i-th token needs to be predicted and
        the mask will be used as input. This token will be included in the loss.
        If target_mask[i] == 0, then the token (or [SEP], [CLS]) will be used as
        input. This token will not be included in the loss.
      tokens: int32 Tensor of shape [seq_length].
      masked_tokens: int32 Tensor of shape [seq_length].
    """
    factorization_length = tf.shape(inputs)[0]
    # Generate permutation indices
    index = tf.range(factorization_length, dtype=tf.int32)
    index = tf.transpose(tf.reshape(index, [-1, self._permutation_size]))
    index = tf.random.shuffle(index)
    index = tf.reshape(tf.transpose(index), [-1])

    input_mask = tf.cast(input_mask, tf.bool)

    # non-functional tokens
    non_func_tokens = tf.logical_not(
        tf.logical_or(
            tf.equal(inputs, self._sep_id), tf.equal(inputs, self._cls_id)))
    masked_tokens = tf.logical_and(input_mask, non_func_tokens)
    non_masked_or_func_tokens = tf.logical_not(masked_tokens)

    smallest_index = -2 * tf.ones([factorization_length], dtype=tf.int32)

    # Similar to BERT, randomly leak some masked tokens
    if self._leak_ratio > 0:
      leak_tokens = tf.logical_and(
          masked_tokens,
          tf.random.uniform([factorization_length], maxval=1.0) <
          self._leak_ratio)
      can_attend_self = tf.logical_or(non_masked_or_func_tokens, leak_tokens)
    else:
      can_attend_self = non_masked_or_func_tokens
    to_index = tf.where(can_attend_self, smallest_index, index)
    from_index = tf.where(can_attend_self, to_index + 1, to_index)

    # For masked tokens, can attend if i > j
    # For context tokens, always can attend each other
    can_attend = from_index[:, None] > to_index[None, :]

    perm_mask = tf.cast(can_attend, tf.int32)

    # Only masked tokens are included in the loss
    target_mask = tf.cast(masked_tokens, tf.int32)

    return perm_mask, target_mask, inputs, masked_tokens
//  [python]
def _parse_predict_data(self, data):
    """Parses data for prediction."""
    # Gets original image and its size.
    image = data['image']
    image_shape = tf.shape(input=image)[0:2]

    # Normalizes image with mean and std pixel values.
    image = input_utils.normalize_image(image)

    # Resizes and crops image.
    image, image_info = input_utils.resize_and_crop_image(
        image,
        self._output_size,
        padded_size=input_utils.compute_padded_size(self._output_size,
                                                    2**self._max_level),
        aug_scale_min=1.0,
        aug_scale_max=1.0)
    image_height, image_width, _ = image.get_shape().as_list()

    # If bfloat16 is used, casts input image to tf.bfloat16.
    if self._use_bfloat16:
      image = tf.cast(image, dtype=tf.bfloat16)

    # Compute Anchor boxes.
    input_anchor = anchor.Anchor(self._min_level, self._max_level,
                                 self._num_scales, self._aspect_ratios,
                                 self._anchor_size, (image_height, image_width))

    labels = {
        'anchor_boxes': input_anchor.multilevel_boxes,
        'image_info': image_info,
    }
    # If mode is PREDICT_WITH_GT, returns groundtruths and training targets
    # in labels.
    if self._mode == ModeKeys.PREDICT_WITH_GT:
      # Converts boxes from normalized coordinates to pixel coordinates.
      boxes = box_utils.denormalize_boxes(data['groundtruth_boxes'],
                                          image_shape)
      groundtruths = {
          'source_id': data['source_id'],
          'num_detections': tf.shape(data['groundtruth_classes']),
          'boxes': boxes,
          'classes': data['groundtruth_classes'],
          'areas': data['groundtruth_area'],
          'is_crowds': tf.cast(data['groundtruth_is_crowd'], tf.int32),
      }
      groundtruths['source_id'] = process_source_id(groundtruths['source_id'])
      groundtruths = pad_groundtruths_to_fixed_size(groundtruths,
                                                    self._max_num_instances)
      labels['groundtruths'] = groundtruths

      # Computes training objective for evaluation loss.
      classes = data['groundtruth_classes']

      image_scale = image_info[2, :]
      offset = image_info[3, :]
      boxes = input_utils.resize_and_crop_boxes(boxes, image_scale,
                                                image_info[1, :], offset)
      # Filters out ground truth boxes that are all zeros.
      indices = box_utils.get_non_empty_box_indices(boxes)
      boxes = tf.gather(boxes, indices)

      # Assigns anchors.
      anchor_labeler = anchor.AnchorLabeler(input_anchor, self._match_threshold,
                                            self._unmatched_threshold)
      (cls_targets, box_targets, num_positives) = anchor_labeler.label_anchors(
          boxes, tf.cast(tf.expand_dims(classes, axis=1), tf.float32))
      labels['cls_targets'] = cls_targets
      labels['box_targets'] = box_targets
      labels['num_positives'] = num_positives
    return image, labels
//  [python]
def _get_segmentation_labels(self, data: TensorDict,
                               unused_features: TensorDict,
                               labels: NestedTensorDict):
    labels['segmentation_output'] = {
        'gt_word_score': t.cast((data['resized_masks'] > 0), t.float32)
    }
//  [python]
def _get_segmentation_labels(self, data: TensorDict,
                               unused_features: TensorDict,
                               labels: NestedTensorDict):
    labels['segmentation_output'] = {
        'gt_word_score': tf.cast((data['resized_masks'] > 0), tf.float32)
    }
//  [python]
def with_geometric_sampling(self, sample_batch_size,
                              discount):
    """Creates tf dataset with geometric sampling.

    Args:
      sample_batch_size: Batch size for sampling.
      discount: MDP discount.

    Returns:
      TensorFlow dataset.
    """

    sample_weights = discount**tf.cast(self.steps, tf.float32)
    weight_sum = tf.math.cumsum(sample_weights)

    def sample_batch(_):
      values = tf.random.uniform((sample_batch_size,), 0.0,
                                 weight_sum[-1])
      ind = tf.searchsorted(weight_sum, values)
      return (tf.gather(self.states, ind,
                        0), tf.gather(self.actions, ind, 0),
              tf.gather(self.next_states, ind,
                        0), tf.gather(self.rewards, ind, 0),
              tf.gather(self.masks, ind,
                        0), tf.gather(self.weights, ind, 0),
              tf.gather(self.steps, ind, 0))

    return tf.data.experimental.Counter().map(sample_batch).prefetch(100)
//  [python]
def _preprocess(self, data: TensorDict, features: TensorDict,
                  unused_labels: TensorDict):
    """All kinds of preprocessing of the decoded data dict."""
    # (1) Decode the entity_id_mask: a H*W*1 mask, each pixel equals to
    #     (1 + position) of the entity in the GT entity list. The IDs
    #     (which can be larger than 255) are stored in the last two channels.
    data['additional_channels'] = t.cast(data['additional_channels'], t.int32)
    entity_id_mask = (
        data['additional_channels'][:, :, -2:-1] * 256 +
        data['additional_channels'][:, :, -1:])
    data['entity_id_mask'] = entity_id_mask

    # (2) Write image id. Used in evaluation.
    if not self._use_tpu:
      features['source_id'] = data['source_id']

    # (3) Block mask: area without annotation
    data['image'] = _erase(
        data['additional_channels'][:, :, 0],
        data['image'],
        min_val=0.,
        max_val=256.)
//  [python]
def _preprocess(self, data: TensorDict, features: TensorDict,
                  unused_labels: TensorDict):
    """All kinds of preprocessing of the decoded data dict."""
    # (1) Decode the entity_id_mask: a H*W*1 mask, each pixel equals to
    #     (1 + position) of the entity in the GT entity list. The IDs
    #     (which can be larger than 255) are stored in the last two channels.
    data['additional_channels'] = tf.cast(data['additional_channels'], tf.int32)
    entity_id_mask = (
        data['additional_channels'][:, :, -2:-1] * 256 +
        data['additional_channels'][:, :, -1:])
    data['entity_id_mask'] = entity_id_mask

    # (2) Write image id. Used in evaluation.
    if not self._use_tpu:
      features['source_id'] = data['source_id']

    # (3) Block mask: area without annotation
    data['image'] = _erase(
        data['additional_channels'][:, :, 0],
        data['image'],
        min_val=0.,
        max_val=256.)
//  [python]
def _erase(mask: tf.Tensor,
           feature: tf.Tensor,
           min_val: float = 0.,
           max_val: float = 256.) -> tf.Tensor:
  """Erase the feature maps with a mask.

  Erase feature maps with a mask and replace the erased area with uniform random
  noise. The mask can have different size from the feature maps.

  Args:
    mask: an (h, w) binay mask for pixels to erase with. Value 1 represents
      pixels to erase.
    feature: the (H, W, C) feature maps to erase from.
    min_val: The minimum value of random noise.
    max_val: The maximum value of random noise.

  Returns:
      The (H, W, C) feature maps, with pixels in mask replaced with noises. It's
    equal to mask * noise + (1 - mask) * feature.
  """
  h, w, c = utilities.resolve_shape(feature)
  resized_mask = tf.image.resize(
      tf.tile(tf.expand_dims(tf.to_float(mask), -1), (1, 1, c)), (h, w))
  erased = tf.where(
      condition=(resized_mask > 0.5),
      x=tf.to_float(tf.random.uniform((h, w, c), min_val, max_val), feature.dtype),
      y=feature)
  return erased
//  [python]
def _erase(mask: tf.Tensor,
           feature: tf.Tensor,
           min_val: float = 0.,
           max_val: float = 256.) -> tf.Tensor:
  """Erase the feature maps with a mask.

  Erase feature maps with a mask and replace the erased area with uniform random
  noise. The mask can have different size from the feature maps.

  Args:
    mask: an (h, w) binay mask for pixels to erase with. Value 1 represents
      pixels to erase.
    feature: the (H, W, C) feature maps to erase from.
    min_val: The minimum value of random noise.
    max_val: The maximum value of random noise.

  Returns:
      The (H, W, C) feature maps, with pixels in mask replaced with noises. It's
    equal to mask * noise + (1 - mask) * feature.
  """
  h, w, c = utilities.resolve_shape(feature)
  resized_mask = tf.image.resize(
      tf.tile(tf.expand_dims(tf.cast(mask, tf.float32), -1), (1, 1, c)), (h, w))
  erased = tf.where(
      condition=(resized_mask > 0.5),
      x=tf.cast(tf.random.uniform((h, w, c), min_val, max_val), feature.dtype),
      y=feature)
  return erased
//  [python]
def safe_mean(losses):
  """Computes a safe mean of the losses.

  Args:
    losses: `Tensor` whose elements contain individual loss measurements.

  Returns:
    A scalar representing the mean of `losses`. If `num_present` is zero,
      then zero is returned.
  """
  total = tf.reduce_sum(losses)
  num_elements = tf.cast(tf.size(losses), dtype=losses.dtype)
  return tf.math.divide_no_nan(total, num_elements)
//  [python]
def _crop_and_resize(self, data: TensorDict, unused_features: TensorDict,
                       unused_labels: TensorDict):
    """Perform random cropping and resizing."""
    # TODO(longshangbang): resize & translate box as well
    # TODO(longshangbang): resize & translate vertices as well

    # Get cropping target.
    h, w = utilities.resolve_shape(data['image'])[:2]
    left, top, crop_w, crop_h, pad_w, pad_h = self._get_crop_box(
        t.cast(h, t.float32), t.cast(w, t.float32))

    # Crop the image. (Pad the images if the crop box is larger than image.)
    if self._is_training:
      # padding left, top, right, bottom
      pad_left = t.random.uniform([], 0, pad_w + 1, dtype=t.int32)
      pad_top = t.random.uniform([], 0, pad_h + 1, dtype=t.int32)
    else:
      pad_left = 0
      pad_top = 0
    cropped_img = tf.image.crop_to_bounding_box(data['image'], top, left,
                                                crop_h, crop_w)
    padded_img = tf.pad(
        cropped_img,
        [[pad_top, pad_h - pad_top], [pad_left, pad_w - pad_left], [0, 0]],
        constant_values=127)

    # Resize images
    data['resized_image'] = tf.image.resize(
        padded_img, (self._output_dimension, self._output_dimension))
    data['resized_image'] = t.cast(data['resized_image'], t.uint8)

    # Crop the masks
    cropped_masks = tf.image.crop_to_bounding_box(data['entity_id_mask'], top,
                                                  left, crop_h, crop_w)
    padded_masks = tf.pad(
        cropped_masks,
        [[pad_top, pad_h - pad_top], [pad_left, pad_w - pad_left], [0, 0]])

    # Resize masks
    data['resized_masks'] = tf.image.resize(
        padded_masks, (self._mask_dimension, self._mask_dimension),
        method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    data['resized_masks'] = tf.squeeze(data['resized_masks'], -1)
//  [python]
def _crop_and_resize(self, data: TensorDict, unused_features: TensorDict,
                       unused_labels: TensorDict):
    """Perform random cropping and resizing."""
    # TODO(longshangbang): resize & translate box as well
    # TODO(longshangbang): resize & translate vertices as well

    # Get cropping target.
    h, w = utilities.resolve_shape(data['image'])[:2]
    left, top, crop_w, crop_h, pad_w, pad_h = self._get_crop_box(
        tf.cast(h, tf.float32), tf.cast(w, tf.float32))

    # Crop the image. (Pad the images if the crop box is larger than image.)
    if self._is_training:
      # padding left, top, right, bottom
      pad_left = tf.random.uniform([], 0, pad_w + 1, dtype=tf.int32)
      pad_top = tf.random.uniform([], 0, pad_h + 1, dtype=tf.int32)
    else:
      pad_left = 0
      pad_top = 0
    cropped_img = tf.image.crop_to_bounding_box(data['image'], top, left,
                                                crop_h, crop_w)
    padded_img = tf.pad(
        cropped_img,
        [[pad_top, pad_h - pad_top], [pad_left, pad_w - pad_left], [0, 0]],
        constant_values=127)

    # Resize images
    data['resized_image'] = tf.image.resize(
        padded_img, (self._output_dimension, self._output_dimension))
    data['resized_image'] = tf.cast(data['resized_image'], tf.uint8)

    # Crop the masks
    cropped_masks = tf.image.crop_to_bounding_box(data['entity_id_mask'], top,
                                                  left, crop_h, crop_w)
    padded_masks = tf.pad(
        cropped_masks,
        [[pad_top, pad_h - pad_top], [pad_left, pad_w - pad_left], [0, 0]])

    # Resize masks
    data['resized_masks'] = tf.image.resize(
        padded_masks, (self._mask_dimension, self._mask_dimension),
        method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    data['resized_masks'] = tf.squeeze(data['resized_masks'], -1)
//  [python]
def build_outputs(self, inputs, mode):
    # If the input image is transposed (from NHWC to HWCN), we need to revert it
    # back to the original shape before it's used in the computation.
    if self._transpose_input:
      inputs = tf.transpose(inputs, [3, 0, 1, 2])

    backbone_features = self._backbone_fn(
        inputs, is_training=(mode == mode_keys.TRAIN))
    fpn_features = self._fpn_fn(
        backbone_features, is_training=(mode == mode_keys.TRAIN))
    cls_outputs, box_outputs = self._head_fn(
        fpn_features, is_training=(mode == mode_keys.TRAIN))

    if self._use_bfloat16:
      levels = cls_outputs.keys()
      for level in levels:
        cls_outputs[level] = tf.to_float(cls_outputs[level])
        box_outputs[level] = tf.to_float(box_outputs[level])

    model_outputs = {
        'cls_outputs': cls_outputs,
        'box_outputs': box_outputs,
    }
    return model_outputs
//  [python]
def build_outputs(self, inputs, mode):
    # If the input image is transposed (from NHWC to HWCN), we need to revert it
    # back to the original shape before it's used in the computation.
    if self._transpose_input:
      inputs = tf.transpose(inputs, [3, 0, 1, 2])

    backbone_features = self._backbone_fn(
        inputs, is_training=(mode == mode_keys.TRAIN))
    fpn_features = self._fpn_fn(
        backbone_features, is_training=(mode == mode_keys.TRAIN))
    cls_outputs, box_outputs = self._head_fn(
        fpn_features, is_training=(mode == mode_keys.TRAIN))

    if self._use_bfloat16:
      levels = cls_outputs.keys()
      for level in levels:
        cls_outputs[level] = tf.cast(cls_outputs[level], tf.float32)
        box_outputs[level] = tf.cast(box_outputs[level], tf.float32)

    model_outputs = {
        'cls_outputs': cls_outputs,
        'box_outputs': box_outputs,
    }
    return model_outputs
//  [python]
def _get_crop_box(
      self, h: tf.Tensor,
      w: tf.Tensor) -> Tuple[Any, Any, tf.Tensor, tf.Tensor, Any, Any]:
    """Get the cropping box.

    Args:
      h: The height of the image to crop. Should be float type.
      w: The width of the image to crop. Should be float type.

    Returns:
      A tuple representing (left, top, crop_w, crop_h, pad_w, pad_h).
      Then in `self._crop_and_resize`, a crop will be extracted with bounding
      box from top-left corner (left, top) and with size (crop_w, crop_h). This
      crop will then be padded with (pad_w, pad_h) to square sizes.
      The outputs also are re-cast to int32 type.
    """
    if not self._is_training or not self._use_cropping:
      # cast back to integers.
      w = t.cast(w, t.int32)
      h = t.cast(h, t.int32)
      side = tf.maximum(w, h)
      return 0, 0, w, h, side - w, side - h

    # Get box size
    scale = tf.random.uniform([], self._crop_min_scale, self._crop_max_scale)
    max_edge = tf.maximum(w, h)
    long_edge = max_edge * scale

    sqrt_aspect_ratio = tf.math.sqrt(
        tf.random.uniform([], self._crop_min_aspect, self._crop_max_aspect))
    box_h = long_edge / sqrt_aspect_ratio
    box_w = long_edge * sqrt_aspect_ratio

    # Get box location
    left = tf.random.uniform([], 0., tf.maximum(0., w - box_w))
    top = tf.random.uniform([], 0., tf.maximum(0., h - box_h))
    # Get crop & pad
    crop_w = tf.minimum(box_w, w - left)
    crop_h = tf.minimum(box_h, h - top)
    pad_w = box_w - crop_w
    pad_h = box_h - crop_h
    return (t.cast(left, t.int32), t.cast(top, t.int32),
            t.cast(crop_w, t.int32), t.cast(crop_h, t.int32),
            t.cast(pad_w, t.int32), t.cast(pad_h, t.int32))
//  [python]
def _get_crop_box(
      self, h: tf.Tensor,
      w: tf.Tensor) -> Tuple[Any, Any, tf.Tensor, tf.Tensor, Any, Any]:
    """Get the cropping box.

    Args:
      h: The height of the image to crop. Should be float type.
      w: The width of the image to crop. Should be float type.

    Returns:
      A tuple representing (left, top, crop_w, crop_h, pad_w, pad_h).
      Then in `self._crop_and_resize`, a crop will be extracted with bounding
      box from top-left corner (left, top) and with size (crop_w, crop_h). This
      crop will then be padded with (pad_w, pad_h) to square sizes.
      The outputs also are re-cast to int32 type.
    """
    if not self._is_training or not self._use_cropping:
      # cast back to integers.
      w = tf.cast(w, tf.int32)
      h = tf.cast(h, tf.int32)
      side = tf.maximum(w, h)
      return 0, 0, w, h, side - w, side - h

    # Get box size
    scale = tf.random.uniform([], self._crop_min_scale, self._crop_max_scale)
    max_edge = tf.maximum(w, h)
    long_edge = max_edge * scale

    sqrt_aspect_ratio = tf.math.sqrt(
        tf.random.uniform([], self._crop_min_aspect, self._crop_max_aspect))
    box_h = long_edge / sqrt_aspect_ratio
    box_w = long_edge * sqrt_aspect_ratio

    # Get box location
    left = tf.random.uniform([], 0., tf.maximum(0., w - box_w))
    top = tf.random.uniform([], 0., tf.maximum(0., h - box_h))
    # Get crop & pad
    crop_w = tf.minimum(box_w, w - left)
    crop_h = tf.minimum(box_h, h - top)
    pad_w = box_w - crop_w
    pad_h = box_h - crop_h
    return (tf.cast(left, tf.int32), tf.cast(top, tf.int32),
            tf.cast(crop_w, tf.int32), tf.cast(crop_h, tf.int32),
            tf.cast(pad_w, tf.int32), tf.cast(pad_h, tf.int32))
//  [python]
def _create_summary_images(self, labels, prediction):
        ''' Create images of the labels and prediction to show training progress
        '''



        images = []

        # Labels is an unsplit tensor, prediction is a split tensor
        split_labels = [ tf.cast(l, floating_point_format) for l in tf.split(labels,len(prediction) , self._channels_dim)]
        prediction = [ tf.expand_dims(tf.cast(p, floating_point_format), self._channels_dim) for p in prediction ]

        if self.args.data_format == "channels_first":
            split_labels = [ tf.transpose(l, [0, 2, 3, 1]) for l in split_labels]
            prediction   = [ tf.transpose(p, [0, 2, 3, 1]) for p in prediction]


        for p in range(len(split_labels)):


                images.append(
                    tf.compat.v1.summary.image('label/plane_{}'.format(p),
                                 split_labels[p],
                                 max_outputs=1)
                    )

                images.append(
                    tf.compat.v1.summary.image('prediction/plane_{}'.format(p),
                                 prediction[p],
                                 max_outputs=1)
                    )

        return tf.compat.v1.summary.merge(images)
//  [python]
def _compute_adv(self, orig, adv_s, beta: float) -> tf.Tensor:
        cond_greater, cond_less_equal, cond_less = self._get_conditions(
            adv_s, orig, beta
        )
        # lower -- min(x^CF - beta, 0.5)
        # upper -- max(x^CF + beta, -0.5)
        upper = tf.minimum(tf.subtract(adv_s, beta), t.float32(0.5))
        lower = tf.maximum(tf.add(adv_s, beta), t.float32(-0.5))
        assign_adv = (
            tf.multiply(cond_greater, upper)
            + tf.multiply(cond_less_equal, orig)
            + tf.multiply(cond_less, lower)
        )

        return self._compute_with_mode(assign_adv, orig)
//  [python]
def _compute_adv(self, orig, adv_s, beta: float) -> tf.Tensor:
        cond_greater, cond_less_equal, cond_less = self._get_conditions(
            adv_s, orig, beta
        )
        # lower -- min(x^CF - beta, 0.5)
        # upper -- max(x^CF + beta, -0.5)
        upper = tf.minimum(tf.subtract(adv_s, beta), tf.cast(0.5, tf.float32))
        lower = tf.maximum(tf.add(adv_s, beta), tf.cast(-0.5, tf.float32))
        assign_adv = (
            tf.multiply(cond_greater, upper)
            + tf.multiply(cond_less_equal, orig)
            + tf.multiply(cond_less, lower)
        )

        return self._compute_with_mode(assign_adv, orig)
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, _ = inputs

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      if self.task_config.train_data.output_audio:
        outputs = model(features, training=True)
      else:
        outputs = model(features['image'], training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(
          lambda x: tf.cast(x, tf.float32), outputs)

      all_losses, contrastive_metrics = self.build_losses(
          model_outputs=outputs, num_replicas=num_replicas,
          model=model)
      loss = all_losses[self.loss]
      scaled_loss = loss

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(
          optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = all_losses
    if metrics:
      self.process_metrics(metrics, contrastive_metrics)
      logs.update({m.name: m.result() for m in metrics})
    return logs
//  [python]
def _get_conditions(
        self, adv: Union[tf.Tensor, tf.Variable], orig, beta: float = 0.0
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
        # (adv - orig) > beta
        cond_greater = t.cast(
            tf.greater(tf.subtract(adv, orig), beta),
            tf.float32,
        )
        # (_adv - orig) <= beta
        cond_less_equal = t.cast(
            tf.less_equal(tf.abs(tf.subtract(adv, orig)), beta),
            tf.float32,
        )
        # (_adv - orig) < beta
        cond_less = t.cast(
            tf.less(tf.subtract(adv, orig), tf.negative(beta)),
            tf.float32,
        )
        return cond_greater, cond_less_equal, cond_less
//  [python]
def _get_conditions(
        self, adv: Union[tf.Tensor, tf.Variable], orig, beta: float = 0.0
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
        # (adv - orig) > beta
        cond_greater = tf.cast(
            tf.greater(tf.subtract(adv, orig), beta),
            tf.float32,
        )
        # (_adv - orig) <= beta
        cond_less_equal = tf.cast(
            tf.less_equal(tf.abs(tf.subtract(adv, orig)), beta),
            tf.float32,
        )
        # (_adv - orig) < beta
        cond_less = tf.cast(
            tf.less(tf.subtract(adv, orig), tf.negative(beta)),
            tf.float32,
        )
        return cond_greater, cond_less_equal, cond_less
//  [python]
def call(self, inputs, mask):
    encoder_shape = tf.shape(mask)
    mask = tf.to_float(mask)
    batch_size, seq_length = encoder_shape[0], encoder_shape[1]
    # reshape for blocking
    blocked_encoder_mask = tf.reshape(
        mask, (batch_size, seq_length // self._block_size, self._block_size))
    encoder_from_mask = tf.reshape(mask, (batch_size, 1, seq_length, 1))
    encoder_to_mask = tf.reshape(mask, (batch_size, 1, 1, seq_length))

    band_mask = create_band_mask_from_inputs(blocked_encoder_mask,
                                             blocked_encoder_mask)
    return [band_mask, encoder_from_mask, encoder_to_mask, blocked_encoder_mask]
//  [python]
def call(self, inputs, mask):
    encoder_shape = tf.shape(mask)
    mask = tf.cast(mask, inputs.dtype)
    batch_size, seq_length = encoder_shape[0], encoder_shape[1]
    # reshape for blocking
    blocked_encoder_mask = tf.reshape(
        mask, (batch_size, seq_length // self._block_size, self._block_size))
    encoder_from_mask = tf.reshape(mask, (batch_size, 1, seq_length, 1))
    encoder_to_mask = tf.reshape(mask, (batch_size, 1, 1, seq_length))

    band_mask = create_band_mask_from_inputs(blocked_encoder_mask,
                                             blocked_encoder_mask)
    return [band_mask, encoder_from_mask, encoder_to_mask, blocked_encoder_mask]
//  [python]
def float32_softmax(x: tf.Tensor, *args, **kwargs) -> tf.Tensor:
  y = tf.cast(tf.nn.softmax(tf.cast(x, tf.float32), *args, **kwargs), x.dtype)
  return y
//  [python]
def causal_attention_mask(nd, ns, dtype):
    """
    1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
    -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return t.cast(m, dtype)
//  [python]
def causal_attention_mask(nd, ns, dtype):
    """
    1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
    -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)
//  [python]
def create_products_of_givens_rotations(dim, seed):
  r"""Constructs a 2D-tensor which is a product of Givens random rotations.
  Constructs a 2D-tensor of the form G_1 * ... * G_k, where G_i is a Givens
  random rotation. The resulting tensor mimics a matrix taken uniformly at
  random form the orthogonal group.
  Args:
    dim: number of rows/columns of the resulting 2D-tensor.
    seed: random seed.
  Returns:
    The product of Givens random rotations.
  """
  nb_givens_rotations = dim * int(math.ceil(math.log(float(dim))))
  q = np.eye(dim, dim)
  np.random.seed(seed)
  for _ in range(nb_givens_rotations):
    random_angle = math.pi * np.random.uniform()
    random_indices = np.random.choice(dim, 2)
    index_i = min(random_indices[0], random_indices[1])
    index_j = max(random_indices[0], random_indices[1])
    slice_i = q[index_i]
    slice_j = q[index_j]
    new_slice_i = math.cos(random_angle) * slice_i + math.sin(
        random_angle) * slice_j
    new_slice_j = -math.sin(random_angle) * slice_i + math.cos(
        random_angle) * slice_j
    q[index_i] = new_slice_i
    q[index_j] = new_slice_j
  return t.cast(t.constant(q), dtype=t.float32)
//  [python]
def create_products_of_givens_rotations(dim, seed):
  r"""Constructs a 2D-tensor which is a product of Givens random rotations.
  Constructs a 2D-tensor of the form G_1 * ... * G_k, where G_i is a Givens
  random rotation. The resulting tensor mimics a matrix taken uniformly at
  random form the orthogonal group.
  Args:
    dim: number of rows/columns of the resulting 2D-tensor.
    seed: random seed.
  Returns:
    The product of Givens random rotations.
  """
  nb_givens_rotations = dim * int(math.ceil(math.log(float(dim))))
  q = np.eye(dim, dim)
  np.random.seed(seed)
  for _ in range(nb_givens_rotations):
    random_angle = math.pi * np.random.uniform()
    random_indices = np.random.choice(dim, 2)
    index_i = min(random_indices[0], random_indices[1])
    index_j = max(random_indices[0], random_indices[1])
    slice_i = q[index_i]
    slice_j = q[index_j]
    new_slice_i = math.cos(random_angle) * slice_i + math.sin(
        random_angle) * slice_j
    new_slice_j = -math.sin(random_angle) * slice_i + math.cos(
        random_angle) * slice_j
    q[index_i] = new_slice_i
    q[index_j] = new_slice_j
  return tf.cast(tf.constant(q), dtype=tf.float32)
//  [python]
def normalized_pixel_coordinates(
    image_width: int, image_height: int) -> Tuple[tf.Tensor, tf.Tensor]:
  """Computes the normalized pixel coordinates for the specified image size.

  The x-coordinates will range from -1 to 1 left to right.
  The y-coordinates will range from -1 to 1 top to bottom.
  The extrema +-1 will fall onto the exterior pixel boundaries, while the
  coordinates will be evaluated at pixel centers. So, image of width 4 will have
  normalized pixel x-coordinates at [-0.75 -0.25 0.25 0.75], while image of
  width 3 will have them at [-0.667 0 0.667].

  Args:
    image_width: int specifying desired output image width in pixels.
    image_height: int specifying desired output image height in pixels.

  Returns:
    Two float32 tensors with shape [image_height, image_width] containing x- and
    y- coordinates, respecively, for each image pixel.
  """
  width = t.cast(image_width, tf.float32)
  height = t.cast(image_height, tf.float32)
  x_range = (2 * tf.range(width) + 1) / width - 1
  y_range = (2 * tf.range(height) + 1) / height - 1
  x_coords, y_coords = tf.meshgrid(x_range, y_range)
  return x_coords, y_coords
//  [python]
def normalized_pixel_coordinates(
    image_width: int, image_height: int) -> Tuple[tf.Tensor, tf.Tensor]:
  """Computes the normalized pixel coordinates for the specified image size.

  The x-coordinates will range from -1 to 1 left to right.
  The y-coordinates will range from -1 to 1 top to bottom.
  The extrema +-1 will fall onto the exterior pixel boundaries, while the
  coordinates will be evaluated at pixel centers. So, image of width 4 will have
  normalized pixel x-coordinates at [-0.75 -0.25 0.25 0.75], while image of
  width 3 will have them at [-0.667 0 0.667].

  Args:
    image_width: int specifying desired output image width in pixels.
    image_height: int specifying desired output image height in pixels.

  Returns:
    Two float32 tensors with shape [image_height, image_width] containing x- and
    y- coordinates, respecively, for each image pixel.
  """
  width = tf.cast(image_width, tf.float32)
  height = tf.cast(image_height, tf.float32)
  x_range = (2 * tf.range(width) + 1) / width - 1
  y_range = (2 * tf.range(height) + 1) / height - 1
  x_coords, y_coords = tf.meshgrid(x_range, y_range)
  return x_coords, y_coords
//  [python]
def loss_model(gt_dict, pred_dict, wt_dict):
        loss = {}
        for k in wt_dict:
            if k == 'pca_verts':
                loss[k] = tf.losses.absolute_difference(gt_dict[k] * gt_dict['garments'][..., np.newaxis],
                                                        pred_dict[k] *
                                                        gt_dict['garments'][...,
                                                                            np.newaxis],
                                                        weights=wt_dict[k])
            elif 'J_2d' in k:
                loss[k] = reprojection(t.cast(gt_dict[k.replace('_2d', '')], tf.float32),
                                       t.cast(pred_dict[k.replace('_2d', '')], tf.float32)) * wt_dict[k]
            else:
                loss[k] = tf.losses.absolute_difference(
                    gt_dict[k], pred_dict[k], weights=wt_dict[k])

        return loss
//  [python]
def loss_model(gt_dict, pred_dict, wt_dict):
        loss = {}
        for k in wt_dict:
            if k == 'pca_verts':
                loss[k] = tf.losses.absolute_difference(gt_dict[k] * gt_dict['garments'][..., np.newaxis],
                                                        pred_dict[k] *
                                                        gt_dict['garments'][...,
                                                                            np.newaxis],
                                                        weights=wt_dict[k])
            elif 'J_2d' in k:
                loss[k] = reprojection(tf.cast(gt_dict[k.replace('_2d', '')], tf.float32),
                                       tf.cast(pred_dict[k.replace('_2d', '')], tf.float32)) * wt_dict[k]
            else:
                loss[k] = tf.losses.absolute_difference(
                    gt_dict[k], pred_dict[k], weights=wt_dict[k])

        return loss
//  [python]
def decode_image(example, feature):
  """Convert image to float32 and normalize from [0, 255] to [0.0, 1.0]."""
  return tf.to_float(feature.decode_example(example)) / 255
//  [python]
def decode_image(example, feature):
  """Convert image to float32 and normalize from [0, 255] to [0.0, 1.0]."""
  return tf.cast(feature.decode_example(example), dtype=tf.float32) / 255
//  [python]
def _create_rpn_loss(self) -> "Tensor":
        """
        Create the partial loss tensor of this attack from RPN losses.

        :return: Attack partial loss tensor.
        """
        import tensorflow as tf

        # Get default graph
        default_graph = tf.get_default_graph()

        # Compute RPN losses
        rpn_iou_threshold = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_iou_threshold")

        # RPN background
        rpn_objectness_predictions_with_background = self.estimator.predictions[
            "rpn_objectness_predictions_with_background"
        ]
        rpn_objectness_predictions_with_background = tf.reshape(
            rpn_objectness_predictions_with_background, (-1, rpn_objectness_predictions_with_background.shape[-1])
        )
        rpn_iou_tensor = default_graph.get_tensor_by_name("Loss/RPNLoss/Compare/IOU/Select:0")
        rpn_iou_tensor = tf.reshape(rpn_iou_tensor, (-1,))
        rpn_target = tf.to_float(tf.greater_equal(rpn_iou_tensor, rpn_iou_threshold))  # pylint: disable=E1123,E1120

        # Compute RPN background loss
        rpn_background_weight = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_background_weight")

        rpn_background_logit = rpn_objectness_predictions_with_background[:, 0]
        rpn_background_loss = rpn_background_logit * rpn_target
        rpn_background_loss = -1 * tf.reduce_sum(rpn_background_loss)
        weight_rpn_background_loss = tf.multiply(
            x=rpn_background_loss, y=rpn_background_weight, name="weight_rpn_background_loss"
        )

        # Compute RPN foreground loss
        rpn_foreground_weight = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_foreground_weight")

        rpn_foreground_logit = rpn_objectness_predictions_with_background[:, 1]
        rpn_foreground_loss = rpn_foreground_logit * rpn_target
        rpn_foreground_loss = tf.reduce_sum(rpn_foreground_loss)
        weight_rpn_foreground_loss = tf.multiply(
            x=rpn_foreground_loss, y=rpn_foreground_weight, name="weight_rpn_foreground_loss"
        )

        # Compute RPN CW loss
        rpn_cw_weight = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_cw_weight")
        rpn_cw_confidence = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_cw_confidence")

        rpn_cw_loss = tf.nn.relu(rpn_foreground_logit - rpn_background_logit + rpn_cw_confidence)
        rpn_cw_loss = rpn_cw_loss * rpn_target
        rpn_cw_loss = tf.reduce_sum(rpn_cw_loss)
        weight_rpn_cw_loss = tf.multiply(x=rpn_cw_loss, y=rpn_cw_weight, name="weight_rpn_cw_loss")

        # Compute partial loss
        partial_loss = tf.add_n(
            [
                weight_rpn_background_loss,
                weight_rpn_foreground_loss,
                weight_rpn_cw_loss,
            ],
            name="partial_rpn_loss",
        )

        return partial_loss