jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _single_permutation_ll(
        self,
        x,
        observed_mask,
        missing_mask=None,
        num_importance_samples=20,
        **map_fn_kwargs,
    ):
        if missing_mask is None:
            missing_mask = tf.broadcast_to(0.0, tf.shape(observed_mask))

        query = (1.0 - observed_mask) * (1.0 - missing_mask)
        observed_mask *= 1.0 - missing_mask

        expanded_x, expanded_observed_mask, selected_features = tf.map_fn(
            _get_autoregressive_ll_batch,
            [x, observed_mask, query],
            fn_output_signature=(
                tf.RaggedTensorSpec((None, x.shape[-1]), x.dtype),
                tf.RaggedTensorSpec((None, x.shape[-1]), observed_mask.dtype),
                tf.RaggedTensorSpec((None, 1), tf.int64),
            ),
            **map_fn_kwargs,
        )

        query_total = tf.math.count_nonzero(query)
        batch_x = expanded_x.merge_dims(0, 1).to_tensor(
            shape=(query_total, self._num_features)
        )
        batch_observed_mask = expanded_observed_mask.merge_dims(0, 1).to_tensor()
        batch_selected_features = tf.expand_dims(selected_features.merge_dims(0, 2), 1)
        batch_selected_features = tf.cast(batch_selected_features, tf.int32)

        outputs = self(
            [batch_x, batch_observed_mask],
            num_importance_samples=num_importance_samples,
            selected_features=batch_selected_features,
        )

        u_count = tf.math.count_nonzero(query, axis=-1)

        energy_ll = tf.RaggedTensor.from_row_lengths(
            tf.squeeze(outputs.energy_ll, -1), u_count
        )
        proposal_ll = tf.RaggedTensor.from_row_lengths(
            tf.squeeze(outputs.proposal_ll, 1), u_count
        )

        energy_ll = tf.reduce_sum(energy_ll, axis=-1)
        proposal_ll = tf.reduce_sum(proposal_ll, axis=-1)

        return energy_ll, proposal_ll
//  [python]
def _instance_autoregressive_sample(self, t):
        x_o, observed_mask, num_samples, num_resampling_samples, use_proposal = t
        u_inds = tf.where((1.0 - observed_mask) == 1)[:, 0]
        u_inds = tf.random.shuffle(u_inds)
        u_inds = t.cast(u_inds, tf.int32)

        x_o = x_o * observed_mask
        curr_x_o = tf.tile(x_o[tf.newaxis, ...], [num_samples, 1])
        curr_observed_mask = tf.tile(observed_mask[tf.newaxis, ...], [num_samples, 1])

        for j in u_inds:
            if use_proposal:
                dim_samples = self._dim_proposal_sample(curr_x_o, curr_observed_mask, j)
            else:
                dim_samples = self._dim_energy_sample(
                    curr_x_o, curr_observed_mask, j, num_resampling_samples
                )

            update_inds = tf.stack(
                [tf.range(num_samples), tf.repeat(j, num_samples)], axis=1
            )

            curr_x_o = tf.tensor_scatter_nd_update(curr_x_o, update_inds, dim_samples)
            curr_observed_mask = tf.tensor_scatter_nd_update(
                curr_observed_mask, update_inds, tf.ones_like(dim_samples)
            )

        return curr_x_o
//  [python]
def _instance_autoregressive_sample(self, t):
        x_o, observed_mask, num_samples, num_resampling_samples, use_proposal = t
        u_inds = tf.where((1.0 - observed_mask) == 1)[:, 0]
        u_inds = tf.random.shuffle(u_inds)
        u_inds = tf.cast(u_inds, tf.int32)

        x_o = x_o * observed_mask
        curr_x_o = tf.tile(x_o[tf.newaxis, ...], [num_samples, 1])
        curr_observed_mask = tf.tile(observed_mask[tf.newaxis, ...], [num_samples, 1])

        for j in u_inds:
            if use_proposal:
                dim_samples = self._dim_proposal_sample(curr_x_o, curr_observed_mask, j)
            else:
                dim_samples = self._dim_energy_sample(
                    curr_x_o, curr_observed_mask, j, num_resampling_samples
                )

            update_inds = tf.stack(
                [tf.range(num_samples), tf.repeat(j, num_samples)], axis=1
            )

            curr_x_o = tf.tensor_scatter_nd_update(curr_x_o, update_inds, dim_samples)
            curr_observed_mask = tf.tensor_scatter_nd_update(
                curr_observed_mask, update_inds, tf.ones_like(dim_samples)
            )

        return curr_x_o
//  [python]
def create_position_ids_from_input_ids(self, input_ids: tf.Tensor) -> tf.Tensor:
        """
        Replace non-padding symbols with their position numbers. Position numbers begin at padding_token_index + 1.
        """
        mask = tf.to_int32(tf.not_equal(input_ids, self.padding_token_index))
        position_ids = tf.cumsum(mask, axis=1) * mask
        position_ids = position_ids + self.padding_token_index
        return position_ids
//  [python]
def create_position_ids_from_input_ids(self, input_ids: tf.Tensor) -> tf.Tensor:
        """
        Replace non-padding symbols with their position numbers. Position numbers begin at padding_token_index + 1.
        """
        mask = tf.cast(tf.not_equal(input_ids, self.padding_token_index), input_ids.dtype)
        position_ids = tf.cumsum(mask, axis=1) * mask
        position_ids = position_ids + self.padding_token_index
        return position_ids
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
    """Compute loss function.

    Args:
      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
      target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
      weights: a float tensor of shape [batch_size, num_anchors]

    Returns:
      loss: a float tensor of shape [batch_size, num_anchors, num_classes]
        representing the value of the loss function.
    """
    if self._bootstrap_type == 'soft':
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
    else:
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.to_float(
              tf.sigmoid(prediction_tensor) > 0.5)
    per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
        labels=bootstrap_target_tensor, logits=prediction_tensor))
    return per_entry_cross_ent * tf.expand_dims(weights, 2)
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
    """Compute loss function.

    Args:
      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
      target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
      weights: a float tensor of shape [batch_size, num_anchors]

    Returns:
      loss: a float tensor of shape [batch_size, num_anchors, num_classes]
        representing the value of the loss function.
    """
    if self._bootstrap_type == 'soft':
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
    else:
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.cast(
              tf.sigmoid(prediction_tensor) > 0.5, tf.float32)
    per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
        labels=bootstrap_target_tensor, logits=prediction_tensor))
    return per_entry_cross_ent * tf.expand_dims(weights, 2)
//  [python]
def _cal_pos_emb(
        self,
        dense_layer: tf.keras.layers.Dense,
        position_ids: tf.Tensor,
        num_buckets: int,
        max_distance: int,
    ):
        rel_pos_matrix = tf.expand_dims(position_ids, axis=-2) - tf.expand_dims(position_ids, axis=-1)
        rel_pos = self.relative_position_bucket(rel_pos_matrix, num_buckets, max_distance)
        rel_pos_one_hot = tf.one_hot(rel_pos, depth=num_buckets, dtype=self.compute_dtype)
        embedding = dense_layer(rel_pos_one_hot)
        # batch_size, seq_length, seq_length, num_heads --> batch_size, num_heads, seq_length, seq_length
        embedding = tf.transpose(embedding, [0, 3, 1, 2])
        embedding = tf.cast(embedding, dtype=self.compute_dtype)
        return embedding
//  [python]
def box_refinement_graph(box, gt_box):
    """Compute refinement needed to transform box to gt_box.
    box and gt_box are [N, (y1, x1, y2, x2)]
    """
    box = tf.cast(box, tf.float32)
    gt_box = tf.cast(gt_box, tf.float32)

    height = box[:, 2] - box[:, 0]
    width = box[:, 3] - box[:, 1]
    center_y = box[:, 0] + 0.5 * height
    center_x = box[:, 1] + 0.5 * width

    gt_height = gt_box[:, 2] - gt_box[:, 0]
    gt_width = gt_box[:, 3] - gt_box[:, 1]
    gt_center_y = gt_box[:, 0] + 0.5 * gt_height
    gt_center_x = gt_box[:, 1] + 0.5 * gt_width

    dy = (gt_center_y - center_y) / height
    dx = (gt_center_x - center_x) / width
    dh = tf.log(gt_height / height)
    dw = tf.log(gt_width / width)

    result = tf.stack([dy, dx, dh, dw], axis=1)
    return result
//  [python]
def update_state(self, targets, prediction, sample_weight=None):
        targets = t.cast(t.squeeze(targets), t.int64)
        prediction = t.cast(t.squeeze(prediction), t.int64)


        values = t.cast(t.equal(prediction,targets), dtype=t.int64)
        not_values = t.cast(t.not_equal(prediction,targets), dtype=t.int64)


        tp = t.where(t.equal(prediction, t.ones_like(prediction) * self.pos_symbol),
                      values, t.zeros_like(values))
        tn = t.where(t.equal(prediction, t.ones_like(prediction) * self.neg_symbol),
                      values, t.zeros_like(values))
        fp = t.where(t.equal(prediction, t.ones_like(prediction) * self.pos_symbol),
                      not_values, t.zeros_like(values))
        fn = t.where(t.equal(prediction, t.ones_like(prediction) * self.neg_symbol),
                      not_values, t.zeros_like(values))

        if sample_weight is not None:
            sample_weight = t.cast(t.squeeze(sample_weight), t.int64)
            tp = t.multiply(tp, sample_weight)
            tn = t.multiply(tn, sample_weight)
            fp = t.multiply(fp, sample_weight)
            fn = t.multiply(fn, sample_weight)

        tp = t.cast(t.reduce_sum(tp), dtype=t.float32)
        tn = t.cast(t.reduce_sum(tn), dtype=t.float32)
        fp = t.cast(t.reduce_sum(fp), dtype=t.float32)
        fn = t.cast(t.reduce_sum(fn), dtype=t.float32)

        self.tp.assign(self.tp + tp)
        self.tn.assign(self.tn + tn)
        self.fp.assign(self.fp + fp)
        self.fn.assign(self.fn + fn)
//  [python]
def update_state(self, targets, prediction, sample_weight=None):
        targets = tf.cast(tf.squeeze(targets), tf.int64)
        prediction = tf.cast(tf.squeeze(prediction), tf.int64)


        values = tf.cast(tf.equal(prediction,targets), dtype=tf.int64)
        not_values = tf.cast(tf.not_equal(prediction,targets), dtype=tf.int64)


        tp = tf.where(tf.equal(prediction, tf.ones_like(prediction) * self.pos_symbol),
                      values, tf.zeros_like(values))
        tn = tf.where(tf.equal(prediction, tf.ones_like(prediction) * self.neg_symbol),
                      values, tf.zeros_like(values))
        fp = tf.where(tf.equal(prediction, tf.ones_like(prediction) * self.pos_symbol),
                      not_values, tf.zeros_like(values))
        fn = tf.where(tf.equal(prediction, tf.ones_like(prediction) * self.neg_symbol),
                      not_values, tf.zeros_like(values))

        if sample_weight is not None:
            sample_weight = tf.cast(tf.squeeze(sample_weight), tf.int64)
            tp = tf.multiply(tp, sample_weight)
            tn = tf.multiply(tn, sample_weight)
            fp = tf.multiply(fp, sample_weight)
            fn = tf.multiply(fn, sample_weight)

        tp = tf.cast(tf.reduce_sum(tp), dtype=tf.float32)
        tn = tf.cast(tf.reduce_sum(tn), dtype=tf.float32)
        fp = tf.cast(tf.reduce_sum(fp), dtype=tf.float32)
        fn = tf.cast(tf.reduce_sum(fn), dtype=tf.float32)

        self.tp.assign(self.tp + tp)
        self.tn.assign(self.tn + tn)
        self.fp.assign(self.fp + fp)
        self.fn.assign(self.fn + fn)
//  [python]
def calculate_visual_bbox(self, batch_size: int, dtype: tf.DType):
        visual_bbox = tf.expand_dims(self.visual_bbox, axis=0)
        visual_bbox = tf.tile(visual_bbox, [batch_size, 1, 1])
        visual_bbox = tf.cast(visual_bbox, dtype=dtype)
        return visual_bbox
//  [python]
def relative_position_bucket(self, relative_positions: tf.Tensor, num_buckets: int, max_distance: int):
        # the negative relative positions are assigned to the interval [0, num_buckets / 2]
        # we deal with this by assigning absolute relative positions to the interval [0, num_buckets / 2]
        # and then offsetting the positive relative positions by num_buckets / 2 at the end
        num_buckets = num_buckets // 2
        buckets = tf.abs(relative_positions)

        # half of the buckets are for exact increments in positions
        max_exact_buckets = num_buckets // 2
        is_small = buckets < max_exact_buckets

        # the other half of the buckets are for logarithmically bigger bins in positions up to max_distance
        buckets_log_ratio = tf.math.log(t.cast(buckets, tf.float32) / max_exact_buckets)
        distance_log_ratio = math.log(max_distance / max_exact_buckets)
        buckets_big_offset = (
            buckets_log_ratio / distance_log_ratio * (num_buckets - max_exact_buckets)
        )  # scale is [0, num_buckets - max_exact_buckets]
        buckets_big = max_exact_buckets + buckets_big_offset  # scale is [max_exact_buckets, num_buckets]
        buckets_big = t.cast(buckets_big, buckets.dtype)
        buckets_big = tf.minimum(buckets_big, num_buckets - 1)

        return (t.cast(relative_positions > 0, buckets.dtype) * num_buckets) + tf.where(
            is_small, buckets, buckets_big
        )
//  [python]
def relative_position_bucket(self, relative_positions: tf.Tensor, num_buckets: int, max_distance: int):
        # the negative relative positions are assigned to the interval [0, num_buckets / 2]
        # we deal with this by assigning absolute relative positions to the interval [0, num_buckets / 2]
        # and then offsetting the positive relative positions by num_buckets / 2 at the end
        num_buckets = num_buckets // 2
        buckets = tf.abs(relative_positions)

        # half of the buckets are for exact increments in positions
        max_exact_buckets = num_buckets // 2
        is_small = buckets < max_exact_buckets

        # the other half of the buckets are for logarithmically bigger bins in positions up to max_distance
        buckets_log_ratio = tf.math.log(tf.cast(buckets, tf.float32) / max_exact_buckets)
        distance_log_ratio = math.log(max_distance / max_exact_buckets)
        buckets_big_offset = (
            buckets_log_ratio / distance_log_ratio * (num_buckets - max_exact_buckets)
        )  # scale is [0, num_buckets - max_exact_buckets]
        buckets_big = max_exact_buckets + buckets_big_offset  # scale is [max_exact_buckets, num_buckets]
        buckets_big = tf.cast(buckets_big, buckets.dtype)
        buckets_big = tf.minimum(buckets_big, num_buckets - 1)

        return (tf.cast(relative_positions > 0, buckets.dtype) * num_buckets) + tf.where(
            is_small, buckets, buckets_big
        )
//  [python]
def call(
        self,
        input_ids: Optional[tf.Tensor] = None,
        bbox: tf.Tensor = None,
        token_type_ids: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        inputs_embeds: Optional[tf.Tensor] = None,
        training: bool = False,
    ) -> tf.Tensor:
        if position_ids is None:
            position_ids = self.create_position_ids(input_ids, inputs_embeds)

        if input_ids is not None:
            input_shape = tf.shape(input_ids)
        else:
            input_shape = tf.shape(inputs_embeds)[:-1]

        if token_type_ids is None:
            token_type_ids = tf.zeros(input_shape, dtype=position_ids.dtype)

        if inputs_embeds is None:
            # Note: tf.gather, on which the embedding layer is based, won't check positive out of bound
            # indices on GPU, returning zeros instead. This is a dangerous silent behavior.
            tf.debugging.assert_less(
                input_ids,
                tf.convert_to_tensor(self.word_embeddings.input_dim, dtype=input_ids.dtype),
                message=(
                    "input_ids must be smaller than the embedding layer's input dimension (got"
                    f" {tf.math.reduce_max(input_ids)} >= {self.word_embeddings.input_dim})"
                ),
            )
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        position_embeddings = self.position_embeddings(position_ids)
        embeddings += position_embeddings

        spatial_position_embeddings = self.calculate_spatial_position_embeddings(bbox)

        embeddings += spatial_position_embeddings

        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings, training=training)
        return embeddings
//  [python]
def call(
        self,
        input_ids: Optional[tf.Tensor] = None,
        bbox: tf.Tensor = None,
        token_type_ids: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        inputs_embeds: Optional[tf.Tensor] = None,
        training: bool = False,
    ) -> tf.Tensor:
        if position_ids is None:
            position_ids = self.create_position_ids(input_ids, inputs_embeds)

        if input_ids is not None:
            input_shape = tf.shape(input_ids)
        else:
            input_shape = tf.shape(inputs_embeds)[:-1]

        if token_type_ids is None:
            token_type_ids = tf.zeros(input_shape, dtype=position_ids.dtype)

        if inputs_embeds is None:
            # Note: tf.gather, on which the embedding layer is based, won't check positive out of bound
            # indices on GPU, returning zeros instead. This is a dangerous silent behavior.
            tf.debugging.assert_less(
                input_ids,
                tf.cast(self.word_embeddings.input_dim, dtype=input_ids.dtype),
                message=(
                    "input_ids must be smaller than the embedding layer's input dimension (got"
                    f" {tf.math.reduce_max(input_ids)} >= {self.word_embeddings.input_dim})"
                ),
            )
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        position_embeddings = self.position_embeddings(position_ids)
        embeddings += position_embeddings

        spatial_position_embeddings = self.calculate_spatial_position_embeddings(bbox)

        embeddings += spatial_position_embeddings

        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings, training=training)
        return embeddings
//  [python]
def switch(condition, then_expression, else_expression):
    """Switches between two operations depending on a scalar value (int or bool).
    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    # Arguments
        condition: scalar tensor.
        then_expression: TensorFlow operation.
        else_expression: TensorFlow operation.
    """
    x_shape = copy.copy(then_expression.get_shape())
    x = tf.cond(t.cast(condition, 'bool'),
                lambda: then_expression,
                lambda: else_expression)
    x.set_shape(x_shape)
    return x
//  [python]
def switch(condition, then_expression, else_expression):
    """Switches between two operations depending on a scalar value (int or bool).
    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    # Arguments
        condition: scalar tensor.
        then_expression: TensorFlow operation.
        else_expression: TensorFlow operation.
    """
    x_shape = copy.copy(then_expression.get_shape())
    x = tf.cond(tf.cast(condition, 'bool'),
                lambda: then_expression,
                lambda: else_expression)
    x.set_shape(x_shape)
    return x
//  [python]
def Recon_Loss(inputs,logits):
    #Calculate Per Sample Reconstruction Loss
    shape_layer_1 = inputs.get_shape().as_list()
    shape_layer_2 = tf.shape(input=inputs)
    recon_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=inputs, logits=logits)
    recon_loss = tf.reshape(recon_loss,shape=[shape_layer_2[0]*shape_layer_2[1],shape_layer_1[2]])
    w=tf.to_float(tf.squeeze(tf.greater(inputs,0),1))
    recon_loss = tf.reduce_mean(input_tensor=w*recon_loss,axis=1)
    return recon_loss
//  [python]
def Recon_Loss(inputs,logits):
    #Calculate Per Sample Reconstruction Loss
    shape_layer_1 = inputs.get_shape().as_list()
    shape_layer_2 = tf.shape(input=inputs)
    recon_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=inputs, logits=logits)
    recon_loss = tf.reshape(recon_loss,shape=[shape_layer_2[0]*shape_layer_2[1],shape_layer_1[2]])
    w=tf.cast(tf.squeeze(tf.greater(inputs,0),1),tf.float32)
    recon_loss = tf.reduce_mean(input_tensor=w*recon_loss,axis=1)
    return recon_loss
//  [python]
def calc_norm_entropy(a):
    return calc_entropy(a) / tf.math.log(tf.to_float(tf.shape(input=a)[0]))
//  [python]
def calc_norm_entropy(a):
    return calc_entropy(a) / tf.math.log(tf.cast(tf.shape(input=a)[0], tf.float32))
//  [python]
def get_extended_attention_mask(self, attention_mask: tf.Tensor) -> tf.Tensor:
        # Adapted from transformers.modelling_utils.ModuleUtilsMixin.get_extended_attention_mask

        n_dims = len(attention_mask.shape)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if n_dims == 3:
            extended_attention_mask = tf.expand_dims(attention_mask, axis=1)
        elif n_dims == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length].
            # Make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length].
            extended_attention_mask = tf.expand_dims(attention_mask, axis=1)  # (batch_size, 1, seq_length)
            extended_attention_mask = tf.expand_dims(extended_attention_mask, axis=1)  # (batch_size, 1, 1, seq_length)
        else:
            raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape}).")

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = tf.to_float(extended_attention_mask)
        extended_attention_mask = (1.0 - extended_attention_mask) * LARGE_NEGATIVE

        return extended_attention_mask
//  [python]
def get_extended_attention_mask(self, attention_mask: tf.Tensor) -> tf.Tensor:
        # Adapted from transformers.modelling_utils.ModuleUtilsMixin.get_extended_attention_mask

        n_dims = len(attention_mask.shape)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if n_dims == 3:
            extended_attention_mask = tf.expand_dims(attention_mask, axis=1)
        elif n_dims == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length].
            # Make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length].
            extended_attention_mask = tf.expand_dims(attention_mask, axis=1)  # (batch_size, 1, seq_length)
            extended_attention_mask = tf.expand_dims(extended_attention_mask, axis=1)  # (batch_size, 1, 1, seq_length)
        else:
            raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape}).")

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = tf.cast(extended_attention_mask, self.compute_dtype)
        extended_attention_mask = (1.0 - extended_attention_mask) * LARGE_NEGATIVE

        return extended_attention_mask
//  [python]
def Get_Gene_Loss(fc,embedding_layer,X_OH):
    upsample1 = tf.compat.v1.layers.dense(fc, 128, tf.nn.relu)
    upsample2 = tf.compat.v1.layers.dense(upsample1, 64, tf.nn.relu)
    upsample3 = tf.compat.v1.layers.dense(upsample2, embedding_layer.shape[1], tf.nn.relu)
    logits = tf.matmul(upsample3, tf.transpose(a=embedding_layer))
    loss = tf.nn.softmax_cross_entropy_with_logits(labels=X_OH, logits=logits)

    predicted = tf.argmax(input=logits,axis=1)
    actual = tf.argmax(input=X_OH,axis=1)
    accuracy = tf.reduce_mean(input_tensor=tf.to_float(tf.equal(predicted,actual)))

    return loss,accuracy
//  [python]
def Get_Gene_Loss(fc,embedding_layer,X_OH):
    upsample1 = tf.compat.v1.layers.dense(fc, 128, tf.nn.relu)
    upsample2 = tf.compat.v1.layers.dense(upsample1, 64, tf.nn.relu)
    upsample3 = tf.compat.v1.layers.dense(upsample2, embedding_layer.shape[1], tf.nn.relu)
    logits = tf.matmul(upsample3, tf.transpose(a=embedding_layer))
    loss = tf.nn.softmax_cross_entropy_with_logits(labels=X_OH, logits=logits)

    predicted = tf.argmax(input=logits,axis=1)
    actual = tf.argmax(input=X_OH,axis=1)
    accuracy = tf.reduce_mean(input_tensor=tf.cast(tf.equal(predicted,actual),tf.float32))

    return loss,accuracy
//  [python]
def Get_HLA_Loss(fc,embedding_layer,X_OH,alpha=1.0):
    upsample1 = tf.compat.v1.layers.dense(fc, 128, tf.nn.relu)
    upsample2 = tf.compat.v1.layers.dense(upsample1, 64, tf.nn.relu)
    upsample3 = tf.compat.v1.layers.dense(upsample2, embedding_layer.shape[1], tf.nn.relu)
    logits = tf.matmul(upsample3, tf.transpose(a=embedding_layer))
    loss = tf.reduce_mean(input_tensor=tf.nn.sigmoid_cross_entropy_with_logits(labels=X_OH, logits=logits),axis=-1)

    predicted = tf.greater(logits,0.9)
    accuracy = tf.reduce_mean(input_tensor=tf.cast(tf.equal(predicted,tf.cast(X_OH,tf.bool)),tf.float32))
    return loss, accuracy
//  [python]
def Convolutional_Features(inputs,reuse=False,prob=0.0,name='Convolutional_Features',kernel=3,net='ae',
                           size_of_net = 'medium',l2_reg=0.0):
    with tf.compat.v1.variable_scope(name,reuse=reuse):
        if size_of_net == 'small':
            units = [12,32,64]
        elif size_of_net == 'medium':
            units = [32,64,128]
        elif size_of_net == 'large':
            units = [64,128,256]
        else:
            units = size_of_net

        for ii,_ in enumerate(units,0):
            if ii == 0:
                conv = tf.compat.v1.layers.conv2d(inputs, units[ii], (1, kernel), 1, padding='same',
                                                  kernel_regularizer=tf.keras.regularizers.l2(l2_reg))
                conv_out = tf.compat.v1.layers.flatten(tf.reduce_max(input_tensor=conv, axis=2))
                indices = tf.squeeze(tf.cast(tf.argmax(input=conv, axis=2), tf.float32), 1)
                conv = tf.nn.leaky_relu(conv)
                conv = tf.compat.v1.layers.dropout(conv, prob)
            else:
                kernel = 3
                conv = tf.compat.v1.layers.conv2d(conv, units[ii], (1, kernel), (1, kernel), padding='same',
                                                  kernel_regularizer=tf.keras.regularizers.l2(l2_reg))
                conv = tf.nn.leaky_relu(conv)
                conv = tf.compat.v1.layers.dropout(conv, prob)

        conv_3 = conv
        conv_3_out = tf.compat.v1.layers.flatten(tf.reduce_max(input_tensor=conv_3,axis=2))

        if net == 'ae':
            return tf.compat.v1.layers.flatten(conv_3),conv_out,indices
        else:
            return conv_3_out,conv_out,indices
//  [python]
def get_mean_embeddings(self, token_embeddings, input_mask):
        """ """
        # cls_embeddings = token_embeddings[:, 0, :]  # 0 is CLS (
)
        # mask PAD tokens
        token_emb_masked = token_embeddings * t.cast(t.expand_dims(input_mask, 2), t.float32)
        total_non_padded_tokens_per_batch = t.cast(t.reduce_sum(input_mask, axis=1), t.float32)
        # Convert to 2D
        total_non_padded_tokens_per_batch = t.expand_dims(total_non_padded_tokens_per_batch, 1)
        mean_embeddings = t.reduce_sum(token_emb_masked, axis=1) / total_non_padded_tokens_per_batch
        return mean_embeddings
//  [python]
def get_mean_embeddings(self, token_embeddings, input_mask):
        """ """
        # cls_embeddings = token_embeddings[:, 0, :]  # 0 is CLS (
)
        # mask PAD tokens
        token_emb_masked = token_embeddings * tf.cast(tf.expand_dims(input_mask, 2), tf.float32)
        total_non_padded_tokens_per_batch = tf.cast(tf.reduce_sum(input_mask, axis=1), tf.float32)
        # Convert to 2D
        total_non_padded_tokens_per_batch = tf.expand_dims(total_non_padded_tokens_per_batch, 1)
        mean_embeddings = tf.reduce_sum(token_emb_masked, axis=1) / total_non_padded_tokens_per_batch
        return mean_embeddings
//  [python]
def build_inputs(self, params, input_context=None):
    """Returns tf.data.Dataset for pretraining."""
    if params.input_path == 'dummy':

      def dummy_data(_):
        dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
        dummy_lm = tf.zeros((1, params.max_predictions_per_seq), dtype=tf.int32)
        return dict(
            input_word_ids=dummy_ids,
            input_mask=dummy_ids,
            input_type_ids=dummy_ids,
            masked_lm_positions=dummy_lm,
            masked_lm_ids=dummy_lm,
            masked_lm_weights=tf.to_float(dummy_lm))

      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return dataset

    return pretrain_dataloader.BertPretrainDataLoader(params).load(
        input_context)
//  [python]
def build_inputs(self, params, input_context=None):
    """Returns tf.data.Dataset for pretraining."""
    if params.input_path == 'dummy':

      def dummy_data(_):
        dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
        dummy_lm = tf.zeros((1, params.max_predictions_per_seq), dtype=tf.int32)
        return dict(
            input_word_ids=dummy_ids,
            input_mask=dummy_ids,
            input_type_ids=dummy_ids,
            masked_lm_positions=dummy_lm,
            masked_lm_ids=dummy_lm,
            masked_lm_weights=tf.cast(dummy_lm, dtype=tf.float32))

      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return dataset

    return pretrain_dataloader.BertPretrainDataLoader(params).load(
        input_context)
//  [python]
def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False):
    """Convert a minibatch of images from uint8 to float32 with configurable dynamic range.
    Can be used as an input transformation for Network.run().
    """
    images = t.cast(images, tf.float32)
    if nhwc_to_nchw:
        images = tf.transpose(images, [0, 3, 1, 2])
    return images * ((drange[1] - drange[0]) / 255) + drange[0]
//  [python]
def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False):
    """Convert a minibatch of images from uint8 to float32 with configurable dynamic range.
    Can be used as an input transformation for Network.run().
    """
    images = tf.cast(images, tf.float32)
    if nhwc_to_nchw:
        images = tf.transpose(images, [0, 3, 1, 2])
    return images * ((drange[1] - drange[0]) / 255) + drange[0]
//  [python]
def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1):
    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
    Can be used as an output transformation for Network.run().
    """
    images = t.cast(images, tf.float32)
    if shrink > 1:
        ksize = [1, 1, shrink, shrink]
        images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
    if nchw_to_nhwc:
        images = tf.transpose(images, [0, 2, 3, 1])
    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return tf.saturate_cast(images, tf.uint8)
//  [python]
def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1):
    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
    Can be used as an output transformation for Network.run().
    """
    images = tf.cast(images, tf.float32)
    if shrink > 1:
        ksize = [1, 1, shrink, shrink]
        images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
    if nchw_to_nhwc:
        images = tf.transpose(images, [0, 2, 3, 1])
    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return tf.saturate_cast(images, tf.uint8)
//  [python]
def _self_suppression(iou, iou_threshold, loop_condition, iou_sum):
  """Bounding-boxes self-suppression loop body.

  Args:
    iou: A float Tensor with shape [1, num_boxes, max_num_instance]: IOUs.
    iou_threshold: A scalar, representing IOU threshold.
    loop_condition: The loop condition returned from last iteration.
    iou_sum: iou_sum_new returned from last iteration.

  Returns:
    iou_suppressed: A float Tensor with shape [1, num_boxes, max_num_instance],
                    IOU after suppression.
    iou_threshold: A scalar, representing IOU threshold.
    loop_condition: Bool Tensor of shape [], the loop condition.
    iou_sum_new: The new IOU sum.
  """
  del loop_condition
  can_suppress_others = t.cast(
      tf.reshape(tf.reduce_max(iou, 1) <= iou_threshold, [1, -1, 1]), iou.dtype)
  iou_suppressed = tf.reshape(
      t.cast(
          tf.reduce_max(can_suppress_others * iou, 1) <= iou_threshold,
          iou.dtype), [1, -1, 1]) * iou
  iou_sum_new = tf.reduce_sum(iou_suppressed, [1, 2])
  return [
      iou_suppressed, iou_threshold,
      tf.reduce_any(iou_sum - iou_sum_new > iou_threshold), iou_sum_new
  ]
//  [python]
def _self_suppression(iou, iou_threshold, loop_condition, iou_sum):
  """Bounding-boxes self-suppression loop body.

  Args:
    iou: A float Tensor with shape [1, num_boxes, max_num_instance]: IOUs.
    iou_threshold: A scalar, representing IOU threshold.
    loop_condition: The loop condition returned from last iteration.
    iou_sum: iou_sum_new returned from last iteration.

  Returns:
    iou_suppressed: A float Tensor with shape [1, num_boxes, max_num_instance],
                    IOU after suppression.
    iou_threshold: A scalar, representing IOU threshold.
    loop_condition: Bool Tensor of shape [], the loop condition.
    iou_sum_new: The new IOU sum.
  """
  del loop_condition
  can_suppress_others = tf.cast(
      tf.reshape(tf.reduce_max(iou, 1) <= iou_threshold, [1, -1, 1]), iou.dtype)
  iou_suppressed = tf.reshape(
      tf.cast(
          tf.reduce_max(can_suppress_others * iou, 1) <= iou_threshold,
          iou.dtype), [1, -1, 1]) * iou
  iou_sum_new = tf.reduce_sum(iou_suppressed, [1, 2])
  return [
      iou_suppressed, iou_threshold,
      tf.reduce_any(iou_sum - iou_sum_new > iou_threshold), iou_sum_new
  ]
//  [python]
def build_losses(self,
                   labels,
                   model_outputs,
                   metrics,
                   aux_losses=None) -> tf.Tensor:
    with tf.name_scope('TeamsPretrainTask/losses'):
      metrics = dict([(metric.name, metric) for metric in metrics])

      # Generator MLM loss.
      lm_prediction_losses = tf_keras.losses.sparse_categorical_crossentropy(
          labels['masked_lm_ids'],
          tf.cast(model_outputs['lm_outputs'], tf.float32),
          from_logits=True)
      lm_label_weights = labels['masked_lm_weights']
      lm_numerator_loss = tf.reduce_sum(lm_prediction_losses * lm_label_weights)
      lm_denominator_loss = tf.reduce_sum(lm_label_weights)
      mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
      metrics['masked_lm_loss'].update_state(mlm_loss)
      weight = self.task_config.model.generator_loss_weight
      total_loss = weight * mlm_loss

      # Discriminator RTD loss.
      rtd_logits = model_outputs['disc_rtd_logits']
      rtd_labels = tf.cast(model_outputs['disc_rtd_label'], tf.float32)
      input_mask = tf.cast(labels['input_mask'], tf.float32)
      rtd_ind_loss = tf.nn.sigmoid_cross_entropy_with_logits(
          logits=rtd_logits, labels=rtd_labels)
      rtd_numerator = tf.reduce_sum(input_mask * rtd_ind_loss)
      rtd_denominator = tf.reduce_sum(input_mask)
      rtd_loss = tf.math.divide_no_nan(rtd_numerator, rtd_denominator)
      metrics['replaced_token_detection_loss'].update_state(rtd_loss)
      weight = self.task_config.model.discriminator_rtd_loss_weight
      total_loss = total_loss + weight * rtd_loss

      # Discriminator MWS loss.
      mws_logits = model_outputs['disc_mws_logits']
      mws_labels = model_outputs['disc_mws_label']
      mws_loss = tf_keras.losses.sparse_categorical_crossentropy(
          mws_labels, mws_logits, from_logits=True)
      mws_numerator_loss = tf.reduce_sum(mws_loss * lm_label_weights)
      mws_denominator_loss = tf.reduce_sum(lm_label_weights)
      mws_loss = tf.math.divide_no_nan(mws_numerator_loss, mws_denominator_loss)
      metrics['multiword_selection_loss'].update_state(mws_loss)
      weight = self.task_config.model.discriminator_mws_loss_weight
      total_loss = total_loss + weight * mws_loss

      if aux_losses:
        total_loss += tf.add_n(aux_losses)

      metrics['total_loss'].update_state(total_loss)
      return total_loss
//  [python]
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
  """Bounding-boxes cross-suppression loop body.

  Args:
    boxes: A float Tensor of shape [1, anchors, 4], representing boxes.
    box_slice: A float Tensor of shape [1, _NMS_TILE_SIZE, 4], the box tile
      returned from last iteration
    iou_threshold: A scalar, representing IOU threshold.
    inner_idx: A scalar, representing inner index.

  Returns:
    boxes: A float Tensor of shape [1, anchors, 4], representing boxes.
    ret_slice: A float Tensor of shape [1, _NMS_TILE_SIZE, 4], the box tile
               after suppression
    iou_threshold: A scalar, representing IOU threshold.
    inner_idx: A scalar, inner index incremented.
  """
  new_slice = tf.slice(boxes, [0, inner_idx * _NMS_TILE_SIZE, 0],
                       [1, _NMS_TILE_SIZE, 4])
  iou = batch_iou(new_slice, box_slice)
  ret_slice = tf.expand_dims(
      tf.to_float(tf.reduce_all(iou < iou_threshold, [1])),
      2) * box_slice
  return boxes, ret_slice, iou_threshold, inner_idx + 1
//  [python]
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
  """Bounding-boxes cross-suppression loop body.

  Args:
    boxes: A float Tensor of shape [1, anchors, 4], representing boxes.
    box_slice: A float Tensor of shape [1, _NMS_TILE_SIZE, 4], the box tile
      returned from last iteration
    iou_threshold: A scalar, representing IOU threshold.
    inner_idx: A scalar, representing inner index.

  Returns:
    boxes: A float Tensor of shape [1, anchors, 4], representing boxes.
    ret_slice: A float Tensor of shape [1, _NMS_TILE_SIZE, 4], the box tile
               after suppression
    iou_threshold: A scalar, representing IOU threshold.
    inner_idx: A scalar, inner index incremented.
  """
  new_slice = tf.slice(boxes, [0, inner_idx * _NMS_TILE_SIZE, 0],
                       [1, _NMS_TILE_SIZE, 4])
  iou = batch_iou(new_slice, box_slice)
  ret_slice = tf.expand_dims(
      tf.cast(tf.reduce_all(iou < iou_threshold, [1]), box_slice.dtype),
      2) * box_slice
  return boxes, ret_slice, iou_threshold, inner_idx + 1
//  [python]
def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None, past_key_values_length: int = 0):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    src_len = shape_list(mask)[1]
    tgt_len = tgt_len if tgt_len is not None else src_len
    one_cst = tf.constant(1.0)
    mask = tf.to_float(mask)
    expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1))

    return (one_cst - expanded_mask) * LARGE_NEGATIVE
//  [python]
def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None, past_key_values_length: int = 0):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    src_len = shape_list(mask)[1]
    tgt_len = tgt_len if tgt_len is not None else src_len
    one_cst = tf.constant(1.0)
    mask = tf.cast(mask, dtype=one_cst.dtype)
    expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1))

    return (one_cst - expanded_mask) * LARGE_NEGATIVE
//  [python]
def _clip_window_prune_boxes(sorted_boxes, clip_window, pad_to_max_output_size,
                             change_coordinate_frame):
  """Prune boxes with zero area.

  Args:
    sorted_boxes: A BoxList containing k detections.
    clip_window: A float32 tensor of the form [y_min, x_min, y_max, x_max]
      representing the window to clip and normalize boxes to before performing
      non-max suppression.
    pad_to_max_output_size: flag indicating whether to pad to max output size or
      not.
    change_coordinate_frame: Whether to normalize coordinates after clipping
      relative to clip_window (this can only be set to True if a clip_window is
      provided).

  Returns:
    sorted_boxes: A BoxList containing k detections after pruning.
    num_valid_nms_boxes_cumulative: Number of valid NMS boxes
  """
  sorted_boxes = box_list_ops.clip_to_window(
      sorted_boxes,
      clip_window,
      filter_nonoverlapping=not pad_to_max_output_size)
  # Set the scores of boxes with zero area to -1 to keep the default
  # behaviour of pruning out zero area boxes.
  sorted_boxes_size = tf.shape(sorted_boxes.get())[0]
  non_zero_box_area = tf.cast(box_list_ops.area(sorted_boxes), tf.bool)
  sorted_boxes_scores = tf.where(
      non_zero_box_area, sorted_boxes.get_field(fields.BoxListFields.scores),
      -1 * tf.ones(sorted_boxes_size))
  sorted_boxes.add_field(fields.BoxListFields.scores, sorted_boxes_scores)
  num_valid_nms_boxes_cumulative = tf.reduce_sum(
      tf.cast(tf.greater_equal(sorted_boxes_scores, 0), tf.int32))
  sorted_boxes = box_list_ops.sort_by_field(sorted_boxes,
                                            fields.BoxListFields.scores)
  if change_coordinate_frame:
    sorted_boxes = box_list_ops.change_coordinate_frame(sorted_boxes,
                                                        clip_window)
  return sorted_boxes, num_valid_nms_boxes_cumulative
//  [python]
def _build_causal_attention_mask(self, batch_size, seq_length, dtype=tf.float32):
        # It is possible with an unspecified sequence length for seq_length to be
        # a runtime value, which is unsupported by tf.constant. Per the TensorFlow
        # docs, tf.fill can handle runtime dynamic shapes:
        # https://www.tensorflow.org/api_docs/python/tf/fill
        diag = tf.to_float(tf.fill((seq_length,), 0.0))

        # set an additive 2D attention mask with all places being masked
        to_mask = tf.to_float(tf.fill((seq_length, seq_length), -10000.0))

        # set diagonal & lower triangular parts to 0 (i.e. the places not to be masked)
        # TIP: think the 2D matrix as the space of (query_seq, key_seq)
        to_mask = tf.linalg.band_part(to_mask, 0, -1)
        # to_mask = tf.linalg.band_part(to_mask, -1, 0)
        to_mask = tf.linalg.set_diag(to_mask, diagonal=diag)

        return tf.broadcast_to(input=to_mask, shape=(batch_size, 1, seq_length, seq_length))
//  [python]
def _build_causal_attention_mask(self, batch_size, seq_length, dtype=tf.float32):
        # It is possible with an unspecified sequence length for seq_length to be
        # a runtime value, which is unsupported by tf.constant. Per the TensorFlow
        # docs, tf.fill can handle runtime dynamic shapes:
        # https://www.tensorflow.org/api_docs/python/tf/fill
        diag = tf.cast(tf.fill((seq_length,), 0.0), dtype)

        # set an additive 2D attention mask with all places being masked
        to_mask = tf.cast(tf.fill((seq_length, seq_length), -10000.0), dtype)

        # set diagonal & lower triangular parts to 0 (i.e. the places not to be masked)
        # TIP: think the 2D matrix as the space of (query_seq, key_seq)
        to_mask = tf.linalg.band_part(to_mask, 0, -1)
        # to_mask = tf.linalg.band_part(to_mask, -1, 0)
        to_mask = tf.linalg.set_diag(to_mask, diagonal=diag)

        return tf.broadcast_to(input=to_mask, shape=(batch_size, 1, seq_length, seq_length))
//  [python]
def _suppression_loop_body(boxes, iou_threshold, output_size, idx):
  """Process boxes in the range [idx*_NMS_TILE_SIZE, (idx+1)*_NMS_TILE_SIZE).

  Args:
    boxes: a tensor with a shape of [1, anchors, 4].
    iou_threshold: a float representing the threshold for deciding whether boxes
      overlap too much with respect to IOU.
    output_size: an int32 tensor of size [1]. Representing the number of
      selected boxes.
    idx: an integer scalar representing induction variable.

  Returns:
    boxes: updated boxes.
    iou_threshold: pass down iou_threshold to the next iteration.
    output_size: the updated output_size.
    idx: the updated induction variable.
  """
  num_tiles = tf.shape(boxes)[1] // _NMS_TILE_SIZE

  # Iterates over tiles that can possibly suppress the current tile.
  box_slice = tf.slice(boxes, [0, idx * _NMS_TILE_SIZE, 0],
                       [1, _NMS_TILE_SIZE, 4])
  _, box_slice, _, _ = tf.while_loop(
      lambda _boxes, _box_slice, _threshold, inner_idx: inner_idx < idx,
      _cross_suppression, [boxes, box_slice, iou_threshold,
                           tf.constant(0)])

  # Iterates over the current tile to compute self-suppression.
  iou = batch_iou(box_slice, box_slice)
  mask = tf.expand_dims(
      tf.reshape(tf.range(_NMS_TILE_SIZE), [1, -1]) > tf.reshape(
          tf.range(_NMS_TILE_SIZE), [-1, 1]), 0)
  iou *= tf.to_float(tf.logical_and(mask, iou >= iou_threshold))
  suppressed_iou, _, _, _ = tf.while_loop(
      lambda _iou, _threshold, loop_condition, _iou_sum: loop_condition,
      _self_suppression,
      [iou, iou_threshold,
       tf.constant(True),
       tf.reduce_sum(iou, [1, 2])])
  suppressed_box = tf.reduce_sum(suppressed_iou, 1) > 0
  box_slice *= tf.expand_dims(1.0 - tf.to_float(suppressed_box), 2)

  # Uses box_slice to update the input boxes.
  mask = tf.reshape(
      tf.to_float(tf.equal(tf.range(num_tiles), idx)), [1, -1, 1, 1])
  boxes = tf.tile(tf.expand_dims(box_slice, [1]),
                  [1, num_tiles, 1, 1]) * mask + tf.reshape(
                      boxes, [1, num_tiles, _NMS_TILE_SIZE, 4]) * (1 - mask)
  boxes = tf.reshape(boxes, [1, -1, 4])

  # Updates output_size.
  output_size += tf.reduce_sum(
      tf.to_int32(tf.reduce_any(box_slice > 0, [2])), [1])
  return boxes, iou_threshold, output_size, idx + 1