jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _token_span_mask(inputs, tgt_len, num_predict, min_num_tokens,
                     max_num_tokens):
  """Sample token spans as prediction targets."""
  mask_alpha = tgt_len / num_predict
  round_to_int = lambda x: tf.cast(tf.round(x), tf.int64)

  # Sample span lengths from a zipf distribution
  span_len_seq = np.arange(min_num_tokens, max_num_tokens + 1)
  probs = np.array([1.0 / (i + 1) for i in span_len_seq])

  probs /= np.sum(probs)
  logits = tf.constant(np.log(probs), dtype=tf.float32)
  span_lens = tf.random.categorical(
      logits=logits[None],
      num_samples=num_predict,
      dtype=tf.int64,
  )[0] + min_num_tokens

  # Sample the ratio [0.0, 1.0) of left context lengths
  span_lens_float = tf.cast(span_lens, tf.float32)
  left_ratio = tf.random.uniform(shape=[num_predict], minval=0.0, maxval=1.0)
  left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)
  left_ctx_len = round_to_int(left_ctx_len)

  # Compute the offset from left start to the right end
  right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

  # Get the actual begin and end indices
  beg_indices = (
      tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
  end_indices = beg_indices + span_lens

  # Remove out of range indices
  valid_idx_mask = end_indices < tgt_len
  beg_indices = tf.boolean_mask(beg_indices, valid_idx_mask)
  end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

  # Shuffle valid indices
  num_valid = tf.cast(tf.shape(beg_indices)[0], tf.int64)
  order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int64))
  beg_indices = tf.gather(beg_indices, order)
  end_indices = tf.gather(end_indices, order)

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)
//  [python]
def _word_span_mask(inputs, tgt_len, num_predict, min_num_words, max_num_words,
                    boundary):
  """Sample whole word spans as prediction targets."""
  # Note: 1.2 is the token-to-word ratio
  mask_alpha = tgt_len / num_predict / 1.2
  round_to_int = lambda x: t.cast(t.round(x), t.int64)

  # Sample span lengths from a zipf distribution
  span_len_seq = np.arange(min_num_words, max_num_words + 1)
  probs = np.array([1.0 / (i + 1) for i in span_len_seq])
  probs /= np.sum(probs)
  logits = t.constant(np.log(probs), dtype=t.float32)

  # Sample `num_predict` words here: note that this is over sampling
  span_lens = t.random.categorical(
      logits=logits[None],
      num_samples=num_predict,
      dtype=t.int64,
  )[0] + min_num_words

  # Sample the ratio [0.0, 1.0) of left context lengths
  span_lens_float = t.cast(span_lens, t.float32)
  left_ratio = t.random.uniform(shape=[num_predict], minval=0.0, maxval=1.0)
  left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)

  left_ctx_len = round_to_int(left_ctx_len)
  right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

  beg_indices = (
      t.cumsum(left_ctx_len) + t.cumsum(right_offset, exclusive=True))
  end_indices = beg_indices + span_lens

  # Remove out of range indices
  max_boundary_index = t.cast(t.shape(boundary)[0] - 1, t.int64)
  valid_idx_mask = end_indices < max_boundary_index
  beg_indices = t.boolean_mask(beg_indices, valid_idx_mask)
  end_indices = t.boolean_mask(end_indices, valid_idx_mask)

  beg_indices = t.gather(boundary, beg_indices)
  end_indices = t.gather(boundary, end_indices)

  # Shuffle valid indices
  num_valid = t.cast(t.shape(beg_indices)[0], t.int64)
  order = t.random.shuffle(t.range(num_valid, dtype=t.int64))
  beg_indices = t.gather(beg_indices, order)
  end_indices = t.gather(end_indices, order)

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)
//  [python]
def _word_span_mask(inputs, tgt_len, num_predict, min_num_words, max_num_words,
                    boundary):
  """Sample whole word spans as prediction targets."""
  # Note: 1.2 is the token-to-word ratio
  mask_alpha = tgt_len / num_predict / 1.2
  round_to_int = lambda x: tf.cast(tf.round(x), tf.int64)

  # Sample span lengths from a zipf distribution
  span_len_seq = np.arange(min_num_words, max_num_words + 1)
  probs = np.array([1.0 / (i + 1) for i in span_len_seq])
  probs /= np.sum(probs)
  logits = tf.constant(np.log(probs), dtype=tf.float32)

  # Sample `num_predict` words here: note that this is over sampling
  span_lens = tf.random.categorical(
      logits=logits[None],
      num_samples=num_predict,
      dtype=tf.int64,
  )[0] + min_num_words

  # Sample the ratio [0.0, 1.0) of left context lengths
  span_lens_float = tf.cast(span_lens, tf.float32)
  left_ratio = tf.random.uniform(shape=[num_predict], minval=0.0, maxval=1.0)
  left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)

  left_ctx_len = round_to_int(left_ctx_len)
  right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

  beg_indices = (
      tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
  end_indices = beg_indices + span_lens

  # Remove out of range indices
  max_boundary_index = tf.cast(tf.shape(boundary)[0] - 1, tf.int64)
  valid_idx_mask = end_indices < max_boundary_index
  beg_indices = tf.boolean_mask(beg_indices, valid_idx_mask)
  end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

  beg_indices = tf.gather(boundary, beg_indices)
  end_indices = tf.gather(boundary, end_indices)

  # Shuffle valid indices
  num_valid = tf.cast(tf.shape(beg_indices)[0], tf.int64)
  order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int64))
  beg_indices = tf.gather(beg_indices, order)
  end_indices = tf.gather(end_indices, order)

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)
//  [python]
def _build_dataset(self, params, vocab_size):
    def dummy_data(_):
      dummy_ids = tf.random.uniform((1, params.seq_length), maxval=vocab_size,
                                    dtype=tf.int32)
      dummy_mask = tf.ones((1, params.seq_length), dtype=tf.int32)
      dummy_type_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_mask,
          input_type_ids=dummy_type_ids,
          masked_lm_positions=dummy_lm,
          masked_lm_ids=dummy_lm,
          masked_lm_weights=tf.to_float(dummy_lm),
          next_sentence_labels=tf.zeros((1, 1), dtype=tf.int32))

    dataset = tf.data.Dataset.range(1)
    dataset = dataset.repeat()
    dataset = dataset.map(
        dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return dataset
//  [python]
def _build_dataset(self, params, vocab_size):
    def dummy_data(_):
      dummy_ids = tf.random.uniform((1, params.seq_length), maxval=vocab_size,
                                    dtype=tf.int32)
      dummy_mask = tf.ones((1, params.seq_length), dtype=tf.int32)
      dummy_type_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_mask,
          input_type_ids=dummy_type_ids,
          masked_lm_positions=dummy_lm,
          masked_lm_ids=dummy_lm,
          masked_lm_weights=tf.cast(dummy_lm, dtype=tf.float32),
          next_sentence_labels=tf.zeros((1, 1), dtype=tf.int32))

    dataset = tf.data.Dataset.range(1)
    dataset = dataset.repeat()
    dataset = dataset.map(
        dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return dataset
//  [python]
def __call__(self, fpn_features, boxes, outer_boxes, classes, is_training):
    """Generate the detection priors from the box detections and FPN features.

    This corresponds to the Fig. 4 of the ShapeMask paper at
    https://arxiv.org/pdf/1904.03239.pdf

    Args:
      fpn_features: a dictionary of FPN features.
      boxes: a float tensor of shape [batch_size, num_instances, 4] representing
        the tight gt boxes from dataloader/detection.
      outer_boxes: a float tensor of shape [batch_size, num_instances, 4]
        representing the loose gt boxes from dataloader/detection.
      classes: a int Tensor of shape [batch_size, num_instances] of instance
        classes.
      is_training: training mode or not.

    Returns:
      instance_features: a float Tensor of shape [batch_size * num_instances,
          mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
          instance feature crop.
      detection_priors: A float Tensor of shape [batch_size * num_instances,
        mask_size, mask_size, 1].
    """
    with tf.name_scope('prior_mask'):
      batch_size, num_instances, _ = boxes.get_shape().as_list()
      outer_boxes = t.cast(outer_boxes, tf.float32)
      boxes = t.cast(boxes, tf.float32)
      instance_features = spatial_transform_ops.multilevel_crop_and_resize(
          fpn_features, outer_boxes, output_size=self._mask_crop_size)
      instance_features = self._shape_prior_fc(instance_features)

      shape_priors = self._get_priors()

      # Get uniform priors for each outer box.
      uniform_priors = tf.ones([
          batch_size, num_instances, self._mask_crop_size, self._mask_crop_size
      ])
      uniform_priors = spatial_transform_ops.crop_mask_in_target_box(
          uniform_priors, boxes, outer_boxes, self._mask_crop_size)

      # Classify shape priors using uniform priors + instance features.
      prior_distribution = self._classify_shape_priors(
          t.cast(instance_features, tf.float32), uniform_priors, classes)

      instance_priors = tf.gather(shape_priors, classes)
      instance_priors *= tf.expand_dims(
          tf.expand_dims(t.cast(prior_distribution, tf.float32), axis=-1),
          axis=-1)
      instance_priors = tf.reduce_sum(instance_priors, axis=2)
      detection_priors = spatial_transform_ops.crop_mask_in_target_box(
          instance_priors, boxes, outer_boxes, self._mask_crop_size)

      return instance_features, detection_priors
//  [python]
def __call__(self, fpn_features, boxes, outer_boxes, classes, is_training):
    """Generate the detection priors from the box detections and FPN features.

    This corresponds to the Fig. 4 of the ShapeMask paper at
    https://arxiv.org/pdf/1904.03239.pdf

    Args:
      fpn_features: a dictionary of FPN features.
      boxes: a float tensor of shape [batch_size, num_instances, 4] representing
        the tight gt boxes from dataloader/detection.
      outer_boxes: a float tensor of shape [batch_size, num_instances, 4]
        representing the loose gt boxes from dataloader/detection.
      classes: a int Tensor of shape [batch_size, num_instances] of instance
        classes.
      is_training: training mode or not.

    Returns:
      instance_features: a float Tensor of shape [batch_size * num_instances,
          mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
          instance feature crop.
      detection_priors: A float Tensor of shape [batch_size * num_instances,
        mask_size, mask_size, 1].
    """
    with tf.name_scope('prior_mask'):
      batch_size, num_instances, _ = boxes.get_shape().as_list()
      outer_boxes = tf.cast(outer_boxes, tf.float32)
      boxes = tf.cast(boxes, tf.float32)
      instance_features = spatial_transform_ops.multilevel_crop_and_resize(
          fpn_features, outer_boxes, output_size=self._mask_crop_size)
      instance_features = self._shape_prior_fc(instance_features)

      shape_priors = self._get_priors()

      # Get uniform priors for each outer box.
      uniform_priors = tf.ones([
          batch_size, num_instances, self._mask_crop_size, self._mask_crop_size
      ])
      uniform_priors = spatial_transform_ops.crop_mask_in_target_box(
          uniform_priors, boxes, outer_boxes, self._mask_crop_size)

      # Classify shape priors using uniform priors + instance features.
      prior_distribution = self._classify_shape_priors(
          tf.cast(instance_features, tf.float32), uniform_priors, classes)

      instance_priors = tf.gather(shape_priors, classes)
      instance_priors *= tf.expand_dims(
          tf.expand_dims(tf.cast(prior_distribution, tf.float32), axis=-1),
          axis=-1)
      instance_priors = tf.reduce_sum(instance_priors, axis=2)
      detection_priors = spatial_transform_ops.crop_mask_in_target_box(
          instance_priors, boxes, outer_boxes, self._mask_crop_size)

      return instance_features, detection_priors
//  [python]
def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):

        if relative_pos is None:
            q = shape_list(query_layer)[-2]
            relative_pos = build_relative_position(q, shape_list(key_layer)[-2])
        shape_list_pos = shape_list(relative_pos)
        if len(shape_list_pos) == 2:
            relative_pos = tf.expand_dims(tf.expand_dims(relative_pos, 0), 0)
        elif len(shape_list_pos) == 3:
            relative_pos = tf.expand_dims(relative_pos, 1)
        # bxhxqxk
        elif len(shape_list_pos) != 4:
            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {len(shape_list_pos)}")

        att_span = tf.minimum(
                tf.maximum(shape_list(query_layer)[-2], shape_list(key_layer)[-2]), self.max_relative_positions
            )
        att_span = tf.cast(att_span,tf.int64)
        rel_embeddings = tf.expand_dims(
            rel_embeddings[self.max_relative_positions - att_span : self.max_relative_positions + att_span, :], 0
        )

        score = 0

        # content->position
        if "c2p" in self.pos_att_type:
            pos_key_layer = self.pos_proj(rel_embeddings)
            pos_key_layer = self.transpose_for_scores(pos_key_layer)
            c2p_att = tf.matmul(query_layer, tf.transpose(pos_key_layer, [0, 1, 3, 2]))
            c2p_pos = tf.clip_by_value(relative_pos + att_span, 0, att_span * 2 - 1)
            c2p_att = torch_gather(c2p_att, c2p_dynamic_expand(c2p_pos, query_layer, relative_pos), -1)
            score += c2p_att

        # position->content
        if "p2c" in self.pos_att_type:
            pos_query_layer = self.pos_q_proj(rel_embeddings)
            pos_query_layer = self.transpose_for_scores(pos_query_layer)
            pos_query_layer /= tf.math.sqrt(tf.cast(shape_list(pos_query_layer)[-1] * scale_factor, dtype=tf.float32))
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                r_pos = build_relative_position(shape_list(key_layer)[-2], shape_list(key_layer)[-2])
            else:
                r_pos = relative_pos
            p2c_pos = tf.clip_by_value(-r_pos + att_span, 0, att_span * 2 - 1)
            p2c_att = tf.matmul(key_layer, tf.transpose(pos_query_layer, [0, 1, 3, 2]))
            p2c_att = tf.transpose(
                torch_gather(p2c_att, p2c_dynamic_expand(p2c_pos, query_layer, key_layer), -1), [0, 1, 3, 2]
            )
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                pos_index = tf.expand_dims(relative_pos[:, :, :, 0], -1)
                p2c_att = torch_gather(p2c_att, pos_dynamic_expand(pos_index, p2c_att, key_layer), -2)
            score += p2c_att

        return score
//  [python]
def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):

        if relative_pos is None:
            q = shape_list(query_layer)[-2]
            relative_pos = build_relative_position(q, shape_list(key_layer)[-2])
        shape_list_pos = shape_list(relative_pos)
        if len(shape_list_pos) == 2:
            relative_pos = tf.expand_dims(tf.expand_dims(relative_pos, 0), 0)
        elif len(shape_list_pos) == 3:
            relative_pos = tf.expand_dims(relative_pos, 1)
        # bxhxqxk
        elif len(shape_list_pos) != 4:
            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {len(shape_list_pos)}")

        att_span = tf.cast(
            tf.minimum(
                tf.maximum(shape_list(query_layer)[-2], shape_list(key_layer)[-2]), self.max_relative_positions
            ),
            tf.int64,
        )
        rel_embeddings = tf.expand_dims(
            rel_embeddings[self.max_relative_positions - att_span : self.max_relative_positions + att_span, :], 0
        )

        score = 0

        # content->position
        if "c2p" in self.pos_att_type:
            pos_key_layer = self.pos_proj(rel_embeddings)
            pos_key_layer = self.transpose_for_scores(pos_key_layer)
            c2p_att = tf.matmul(query_layer, tf.transpose(pos_key_layer, [0, 1, 3, 2]))
            c2p_pos = tf.clip_by_value(relative_pos + att_span, 0, att_span * 2 - 1)
            c2p_att = torch_gather(c2p_att, c2p_dynamic_expand(c2p_pos, query_layer, relative_pos), -1)
            score += c2p_att

        # position->content
        if "p2c" in self.pos_att_type:
            pos_query_layer = self.pos_q_proj(rel_embeddings)
            pos_query_layer = self.transpose_for_scores(pos_query_layer)
            pos_query_layer /= tf.math.sqrt(tf.cast(shape_list(pos_query_layer)[-1] * scale_factor, dtype=tf.float32))
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                r_pos = build_relative_position(shape_list(key_layer)[-2], shape_list(key_layer)[-2])
            else:
                r_pos = relative_pos
            p2c_pos = tf.clip_by_value(-r_pos + att_span, 0, att_span * 2 - 1)
            p2c_att = tf.matmul(key_layer, tf.transpose(pos_query_layer, [0, 1, 3, 2]))
            p2c_att = tf.transpose(
                torch_gather(p2c_att, p2c_dynamic_expand(p2c_pos, query_layer, key_layer), -1), [0, 1, 3, 2]
            )
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                pos_index = tf.expand_dims(relative_pos[:, :, :, 0], -1)
                p2c_att = torch_gather(p2c_att, pos_dynamic_expand(pos_index, p2c_att, key_layer), -2)
            score += p2c_att

        return score
//  [python]
def _convert_features(
      self,
      ds,
      task_feature_lengths,
  ):
    """Convert the input dataset to an output dataset to be fed to the model.

    The conversion process involves three steps
    1. The pair of sequences and any of the requested metadata are fetched from
       `example`.
    2. The state sequence is encoded as an alignment path.
    3. The sequences and alignment path are padded. Packing is not yet
       supported.

    Args:
      ds: an input tf.data.Dataset to be converted.
      task_feature_lengths: a mapping from a feature to its length.

    Returns:
      ds: the converted dataset.
    """

    @seqio.map_over_dataset
    def convert_example(
        example,
    ):
      out = {'sequence_x': example['sequence_x'],
             'sequence_y': example['sequence_y']}
      if self.include_metadata_features:
        for k in self.METADATA_FEATURES:
          out[k] = tf.reshape(example[k], [-1])

      states = example['states']
      ali_start_x, ali_start_y = example['ali_start_x'], example['ali_start_y']

      matches = t.cast(states[1:] == self._match, self.dtype)
      gaps_in_x = t.cast(states[1:] == self._gap_in_x, self.dtype)
      gaps_in_y = t.cast(states[1:] == self._gap_in_y, self.dtype)
      padding_mask = states[1:] != self._pad_id

      out['align_path_x'] = tf.where(
          padding_mask, ali_start_x - 1 + tf.cumsum(matches + gaps_in_y), 0)
      out['align_path_y'] = tf.where(
          padding_mask, ali_start_y - 1 + tf.cumsum(matches + gaps_in_x), 0)
      out['align_path_s'] = tf.where(
          padding_mask,
          tf.gather(
              params=self._trans_encoder,
              indices=self._hash_fn(states[:-1], states[1:])),
          0)

      return out

    ds = self._pack_or_pad(ds, task_feature_lengths)
    return convert_example(ds)
//  [python]
def _convert_features(
      self,
      ds,
      task_feature_lengths,
  ):
    """Convert the input dataset to an output dataset to be fed to the model.

    The conversion process involves three steps
    1. The pair of sequences and any of the requested metadata are fetched from
       `example`.
    2. The state sequence is encoded as an alignment path.
    3. The sequences and alignment path are padded. Packing is not yet
       supported.

    Args:
      ds: an input tf.data.Dataset to be converted.
      task_feature_lengths: a mapping from a feature to its length.

    Returns:
      ds: the converted dataset.
    """

    @seqio.map_over_dataset
    def convert_example(
        example,
    ):
      out = {'sequence_x': example['sequence_x'],
             'sequence_y': example['sequence_y']}
      if self.include_metadata_features:
        for k in self.METADATA_FEATURES:
          out[k] = tf.reshape(example[k], [-1])

      states = example['states']
      ali_start_x, ali_start_y = example['ali_start_x'], example['ali_start_y']

      matches = tf.cast(states[1:] == self._match, self.dtype)
      gaps_in_x = tf.cast(states[1:] == self._gap_in_x, self.dtype)
      gaps_in_y = tf.cast(states[1:] == self._gap_in_y, self.dtype)
      padding_mask = states[1:] != self._pad_id

      out['align_path_x'] = tf.where(
          padding_mask, ali_start_x - 1 + tf.cumsum(matches + gaps_in_y), 0)
      out['align_path_y'] = tf.where(
          padding_mask, ali_start_y - 1 + tf.cumsum(matches + gaps_in_x), 0)
      out['align_path_s'] = tf.where(
          padding_mask,
          tf.gather(
              params=self._trans_encoder,
              indices=self._hash_fn(states[:-1], states[1:])),
          0)

      return out

    ds = self._pack_or_pad(ds, task_feature_lengths)
    return convert_example(ds)
//  [python]
def _convert_features(
      self,
      ds,
      task_feature_lengths,
  ):
    """Convert the input dataset to an output dataset to be fed to the model.

    The conversion process involves three steps
    1. The pair of sequences and any of the requested metadata are fetched from
       `example`.
    2. The homology label is fetched from `example` and, if binarized if
       `self.fine_grained_labels` is `True`.
    3. The sequences are padded. Packing is not yet supported.

    Args:
      ds: an input tf.data.Dataset to be converted.
      task_feature_lengths: a mapping from a feature to its length.

    Returns:
      ds: the converted dataset.
    """

    @seqio.map_over_dataset
    def convert_example(
        example,
    ):
      homology_label = tf.reshape(example['homology_label'], [-1])
      if not self.fine_grained_labels:
        homology_label = tf.cast(homology_label > 0, homology_label.dtype)
      out = {'sequence_x': example['sequence_x'],
             'sequence_y': example['sequence_y'],
             'homology_label': homology_label}
      if self.include_metadata_features:
        for k in self.METADATA_FEATURES:
          out[k] = tf.reshape(example[k], [-1])

      return out

    ds = self._pack_or_pad(ds, task_feature_lengths)
    return convert_example(ds)
//  [python]
def __call__(self, features, mask_logits, classes, is_training):
    """Generate instance masks from FPN features and detection priors.

    This corresponds to the Fig. 5-6 of the ShapeMask paper at
    https://arxiv.org/pdf/1904.03239.pdf

    Args:
      features: a float Tensor of shape [batch_size, num_instances,
        mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
        instance feature crop.
      mask_logits: a float Tensor of shape [batch_size, num_instances,
        mask_crop_size, mask_crop_size] indicating predicted mask logits.
      classes: a int Tensor of shape [batch_size, num_instances] of instance
        classes.
      is_training: a bool indicating whether in training mode.

    Returns:
      mask_outputs: instance mask prediction as a float Tensor of shape
        [batch_size, num_instances, mask_size, mask_size].
    """
    # Extract the foreground mean features
    # with tf.variable_scope('fine_mask', reuse=tf.AUTO_REUSE):
    with tf.name_scope('fine_mask'):
      mask_probs = tf.nn.sigmoid(mask_logits)
      # Compute instance embedding for hard average.
      binary_mask = tf.to_int32(tf.greater(mask_probs, 0.5))
      instance_embedding = tf.reduce_sum(
          features * tf.expand_dims(tf.cast(binary_mask, features.dtype), axis=-1), axis=(2, 3))
      instance_embedding /= tf.expand_dims(
          tf.reduce_sum(tf.cast(binary_mask, features.dtype), axis=(2, 3)) + 1e-20, axis=-1)
      # Take the difference between crop features and mean instance features.
      features -= tf.expand_dims(
          tf.expand_dims(instance_embedding, axis=2), axis=2)

      features += self._fine_mask_fc(tf.expand_dims(mask_probs, axis=-1))

      # Decoder to generate upsampled segmentation mask.
      mask_logits = self.decoder_net(features, is_training)
      if self._use_category_for_mask:
        mask_logits = tf.transpose(mask_logits, [0, 1, 4, 2, 3])
        mask_logits = tf.gather(
            mask_logits, tf.expand_dims(classes, -1), batch_dims=2)
        mask_logits = tf.squeeze(mask_logits, axis=2)
      else:
        mask_logits = mask_logits[..., 0]

    return mask_logits
//  [python]
def __call__(self, features, mask_logits, classes, is_training):
    """Generate instance masks from FPN features and detection priors.

    This corresponds to the Fig. 5-6 of the ShapeMask paper at
    https://arxiv.org/pdf/1904.03239.pdf

    Args:
      features: a float Tensor of shape [batch_size, num_instances,
        mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
        instance feature crop.
      mask_logits: a float Tensor of shape [batch_size, num_instances,
        mask_crop_size, mask_crop_size] indicating predicted mask logits.
      classes: a int Tensor of shape [batch_size, num_instances] of instance
        classes.
      is_training: a bool indicating whether in training mode.

    Returns:
      mask_outputs: instance mask prediction as a float Tensor of shape
        [batch_size, num_instances, mask_size, mask_size].
    """
    # Extract the foreground mean features
    # with tf.variable_scope('fine_mask', reuse=tf.AUTO_REUSE):
    with tf.name_scope('fine_mask'):
      mask_probs = tf.nn.sigmoid(mask_logits)
      # Compute instance embedding for hard average.
      binary_mask = tf.cast(tf.greater(mask_probs, 0.5), features.dtype)
      instance_embedding = tf.reduce_sum(
          features * tf.expand_dims(binary_mask, axis=-1), axis=(2, 3))
      instance_embedding /= tf.expand_dims(
          tf.reduce_sum(binary_mask, axis=(2, 3)) + 1e-20, axis=-1)
      # Take the difference between crop features and mean instance features.
      features -= tf.expand_dims(
          tf.expand_dims(instance_embedding, axis=2), axis=2)

      features += self._fine_mask_fc(tf.expand_dims(mask_probs, axis=-1))

      # Decoder to generate upsampled segmentation mask.
      mask_logits = self.decoder_net(features, is_training)
      if self._use_category_for_mask:
        mask_logits = tf.transpose(mask_logits, [0, 1, 4, 2, 3])
        mask_logits = tf.gather(
            mask_logits, tf.expand_dims(classes, -1), batch_dims=2)
        mask_logits = tf.squeeze(mask_logits, axis=2)
      else:
        mask_logits = mask_logits[..., 0]

    return mask_logits
//  [python]
def call(self,
           inputs,
           mask,
           inverse_normalizer,
           step=None,
           beam_indices=None,
           cache=None,
           attn_mask=None):
    self._assert_rank_and_type(inputs, 3)
    self._assert_rank_and_type(mask, 3)
    assert inputs.get_shape().as_list()[-1] == self.model_dimension

    layer_out = self.dense_layers(inputs, mask, inverse_normalizer)

    # TFLite mode is handled with a custom op.
    if self.parameters.mode == base_layers.TFLITE:
      assert beam_indices is not None
      assert step is not None
      layer_out = tf_custom_ops_py.uniform_causal_attn(
          layer_out, step, beam_indices, self.model_dimension, self.beam_size)
    else:
      # Cache is used for TF Predict and Eval modes.
      if cache is None:
        attention_matrix = self.get_uniform_attention(attn_mask)
        layer_out = tf.matmul(attention_matrix, layer_out)
      else:
        assert self.parameters.mode in [base_layers.PREDICT, base_layers.EVAL]
        assert step is not None
        cache['uniform_avg'] = layer_out + cache['uniform_avg']
        layer_out = cache['uniform_avg'] / t.cast(step, dtype=t.float32)
    return self.qoutput(layer_out)
//  [python]
def call(self,
           inputs,
           mask,
           inverse_normalizer,
           step=None,
           beam_indices=None,
           cache=None,
           attn_mask=None):
    self._assert_rank_and_type(inputs, 3)
    self._assert_rank_and_type(mask, 3)
    assert inputs.get_shape().as_list()[-1] == self.model_dimension

    layer_out = self.dense_layers(inputs, mask, inverse_normalizer)

    # TFLite mode is handled with a custom op.
    if self.parameters.mode == base_layers.TFLITE:
      assert beam_indices is not None
      assert step is not None
      layer_out = tf_custom_ops_py.uniform_causal_attn(
          layer_out, step, beam_indices, self.model_dimension, self.beam_size)
    else:
      # Cache is used for TF Predict and Eval modes.
      if cache is None:
        attention_matrix = self.get_uniform_attention(attn_mask)
        layer_out = tf.matmul(attention_matrix, layer_out)
      else:
        assert self.parameters.mode in [base_layers.PREDICT, base_layers.EVAL]
        assert step is not None
        cache['uniform_avg'] = layer_out + cache['uniform_avg']
        layer_out = cache['uniform_avg'] / tf.cast(step, dtype=tf.float32)
    return self.qoutput(layer_out)
//  [python]
def call(self, input_indices, sequence_length):
    mask_rank2 = tf.sequence_mask(
        sequence_length, tf.shape(input_indices)[1], dtype=tf.float32)
    mask_rank3 = tf.expand_dims(mask_rank2, axis=2)
    inverse_normalizer = tf.math.reciprocal(tf.reduce_sum(mask_rank3))
    if self.parameters.mode in [base_layers.PREDICT, base_layers.TFLITE]:
      sequence_length = tf.reduce_sum(input_indices + 1 - input_indices)
      pos_indices = tf.range(sequence_length, dtype=tf.int32)
      pos_indices = tf.reshape(pos_indices, [1, -1])
    else:
      pos_indices = tf.cumsum(mask_rank2, axis=1, exclusive=True)
      pos_indices = tf.to_int32(pos_indices)

    input_values = self.embedding(input_indices)
    pos_values = self.positional_embedding(pos_indices)
    inputs = self.qact(self.ln(input_values + pos_values))
    attn_mask = tf.matmul(mask_rank3, tf.transpose(mask_rank3, [0, 2, 1]))
    if self.parameters.mode not in [base_layers.PREDICT, base_layers.TFLITE]:
      inputs = inputs * mask_rank3
    for layer in self.layers:
      outputs = layer(inputs, mask_rank3, inverse_normalizer, attn_mask)
      inputs = outputs
    if self.parameters.mode not in [base_layers.PREDICT, base_layers.TFLITE]:
      outputs = outputs * mask_rank3
    return outputs
//  [python]
def call(self, input_indices, sequence_length):
    mask_rank2 = tf.sequence_mask(
        sequence_length, tf.shape(input_indices)[1], dtype=tf.float32)
    mask_rank3 = tf.expand_dims(mask_rank2, axis=2)
    inverse_normalizer = tf.math.reciprocal(tf.reduce_sum(mask_rank3))
    if self.parameters.mode in [base_layers.PREDICT, base_layers.TFLITE]:
      sequence_length = tf.reduce_sum(input_indices + 1 - input_indices)
      pos_indices = tf.range(sequence_length, dtype=tf.int32)
      pos_indices = tf.reshape(pos_indices, [1, -1])
    else:
      pos_indices = tf.cumsum(mask_rank2, axis=1, exclusive=True)
      pos_indices = tf.cast(pos_indices, dtype=tf.int32)

    input_values = self.embedding(input_indices)
    pos_values = self.positional_embedding(pos_indices)
    inputs = self.qact(self.ln(input_values + pos_values))
    attn_mask = tf.matmul(mask_rank3, tf.transpose(mask_rank3, [0, 2, 1]))
    if self.parameters.mode not in [base_layers.PREDICT, base_layers.TFLITE]:
      inputs = inputs * mask_rank3
    for layer in self.layers:
      outputs = layer(inputs, mask_rank3, inverse_normalizer, attn_mask)
      inputs = outputs
    if self.parameters.mode not in [base_layers.PREDICT, base_layers.TFLITE]:
      outputs = outputs * mask_rank3
    return outputs
//  [python]
def accuracy_fn(inference_fn, inputs, labels):
    prediction = tf.nn.softmax(inference_fn(inputs))
    correct_pred = tf.equal(tf.argmax(prediction, 1), labels)
    return tf.reduce_mean(tf.to_float(correct_pred))
//  [python]
def accuracy_fn(inference_fn, inputs, labels):
    prediction = tf.nn.softmax(inference_fn(inputs))
    correct_pred = tf.equal(tf.argmax(prediction, 1), labels)
    return tf.reduce_mean(tf.cast(correct_pred, tf.float32))
//  [python]
def points_inside_cube(self,
                           points,
                           bottom_centers,
                           rotation,
                           cube_size=[1.45, 1.55, 4.00]):
        bottom_centers = tf.expand_dims(
                           bottom_centers, axis=2)
        points_centered = points - bottom_centers
        x, y, z = tf.split(points_centered,
                           [1, 1, 1], axis=1)
        xz = tf.concat([x, z], axis=1)
        rot = tf.reshape(rotation, (-1, 1, 1))
        center_to_front_unit_vect = \
            tf.concat([tf.math.cos(rot),
                       -tf.math.sin(rot)], axis=1)
        center_to_side_unit_vect = \
            tf.concat([tf.math.sin(rot),
                       tf.math.cos(rot)], axis=1)
        proj_front = tf.reduce_sum(
                         xz * center_to_front_unit_vect,
                         axis=1, keepdims=True)
        proj_side =  tf.reduce_sum(
                        xz * center_to_side_unit_vect,
                        axis=1, keepdims=True)
        inside_front = tf.to_float(tf.less(tf.abs(proj_front), 
                                   cube_size[2] * 0.5))
        inside_side = tf.to_float(tf.less(tf.abs(proj_side), 
                                  cube_size[1] * 0.5))
        inside_bottom = tf.to_float(tf.less(y, 0.0))
        inside_mask = inside_front * inside_side * \
                      inside_bottom
        
        return inside_mask, proj_front, proj_side
//  [python]
def points_inside_cube(self,
                           points,
                           bottom_centers,
                           rotation,
                           cube_size=[1.45, 1.55, 4.00]):
        bottom_centers = tf.expand_dims(
                           bottom_centers, axis=2)
        points_centered = points - bottom_centers
        x, y, z = tf.split(points_centered,
                           [1, 1, 1], axis=1)
        xz = tf.concat([x, z], axis=1)
        rot = tf.reshape(rotation, (-1, 1, 1))
        center_to_front_unit_vect = \
            tf.concat([tf.math.cos(rot),
                       -tf.math.sin(rot)], axis=1)
        center_to_side_unit_vect = \
            tf.concat([tf.math.sin(rot),
                       tf.math.cos(rot)], axis=1)
        proj_front = tf.reduce_sum(
                         xz * center_to_front_unit_vect,
                         axis=1, keepdims=True)
        proj_side =  tf.reduce_sum(
                        xz * center_to_side_unit_vect,
                        axis=1, keepdims=True)
        inside_front = tf.cast(tf.less(tf.abs(proj_front), 
                                   cube_size[2] * 0.5),
                               tf.float32)
        inside_side = tf.cast(tf.less(tf.abs(proj_side), 
                                  cube_size[1] * 0.5),
                              tf.float32)
        inside_bottom = tf.cast(tf.less(y, 0.0), tf.float32)
        inside_mask = inside_front * inside_side * \
                      inside_bottom
        
        return inside_mask, proj_front, proj_side
//  [python]
def solve_for_constrained_w_row_with_pseudoinv(
    h_matrix: tf.Tensor, s_row: tf.Tensor,
    row_constraint: tf.Tensor) -> tf.Tensor:
  r"""Solves for a single row of W with minimal Frobenius norm.

  Args:
    h_matrix: rank-2 tensor representing the term H in the factorization WH = S.
    s_row: rank-1 tensor representing the term S in the factorization WH = S.
    row_constraint: rank-1 tensor of Booleans representing constraints on the
      matrix W.

  Returns:
    The matrix W above; the unique rank-2 tensor of minimal Frobenius norm
    satisfying WH = S, with all elements corresponding to entries in
    constraint_matrix which are False guaranteed to be zero.
  """
  e_matrix = _construct_e_matrix(row_constraint, h_matrix.shape[0])
  row_soln = tf.squeeze(
      tf.reshape(t.cast(s_row, h_matrix.dtype), [1, -1]) @ tf.linalg.pinv(
          t.cast(e_matrix, h_matrix.dtype) @ h_matrix))
  return row_soln
//  [python]
def solve_for_constrained_w_row_with_pseudoinv(
    h_matrix: tf.Tensor, s_row: tf.Tensor,
    row_constraint: tf.Tensor) -> tf.Tensor:
  r"""Solves for a single row of W with minimal Frobenius norm.

  Args:
    h_matrix: rank-2 tensor representing the term H in the factorization WH = S.
    s_row: rank-1 tensor representing the term S in the factorization WH = S.
    row_constraint: rank-1 tensor of Booleans representing constraints on the
      matrix W.

  Returns:
    The matrix W above; the unique rank-2 tensor of minimal Frobenius norm
    satisfying WH = S, with all elements corresponding to entries in
    constraint_matrix which are False guaranteed to be zero.
  """
  e_matrix = _construct_e_matrix(row_constraint, h_matrix.shape[0])
  row_soln = tf.squeeze(
      tf.reshape(tf.cast(s_row, h_matrix.dtype), [1, -1]) @ tf.linalg.pinv(
          tf.cast(e_matrix, h_matrix.dtype) @ h_matrix))
  return row_soln
//  [python]
def _construct_e_matrix(row_constraint: tf.Tensor,
                        row_dim: tf.Tensor) -> tf.Tensor:
  r"""Constructs projection operation for constrained pseudoinversion.

  The matrix E constructed is one part of the so-called
  'constrained generalized inverse'; the constrained generalized inverse of A
  with respect to the (axis-aligned) constraints on the solution vector
  represented by the boolean `row_constraint` will be `(EA)^\dagger`.

  For a derivation of the expressions used here, see section 3.6 of Campbell
  and Meyer's 'Generalized Inverses of Linear Transformations'

  Args:
    row_constraint: A rank-1 Boolean tensor indicating which elements may be
      nonzero.
    row_dim: Integer tensor representing the intended number of rows of the
      matrix `A` as described above.

  Returns:
    A rank-2 tensor E as described above.
  """
  num_false_elem = tf.reduce_sum(
      tf.ones(shape=row_constraint.shape, dtype=tf.int32) -
      tf.to_int32(row_constraint))
  if num_false_elem == 0:
    e_matrix = tf.eye(row_dim)
  else:
    c_matrix = tf.zeros(shape=[num_false_elem, tf.size(row_constraint)])
    col_idx = tf.size(row_constraint) - num_false_elem
    updates = [1]
    for row_idx in range(num_false_elem):
      indices = [[row_idx, col_idx]]
      c_matrix = tf.tensor_scatter_nd_add(c_matrix, indices, updates)
      col_idx += 1
    e_matrix = tf.eye(row_dim) - tf.linalg.pinv(c_matrix) @ c_matrix
  return e_matrix
//  [python]
def _construct_e_matrix(row_constraint: tf.Tensor,
                        row_dim: tf.Tensor) -> tf.Tensor:
  r"""Constructs projection operation for constrained pseudoinversion.

  The matrix E constructed is one part of the so-called
  'constrained generalized inverse'; the constrained generalized inverse of A
  with respect to the (axis-aligned) constraints on the solution vector
  represented by the boolean `row_constraint` will be `(EA)^\dagger`.

  For a derivation of the expressions used here, see section 3.6 of Campbell
  and Meyer's 'Generalized Inverses of Linear Transformations'

  Args:
    row_constraint: A rank-1 Boolean tensor indicating which elements may be
      nonzero.
    row_dim: Integer tensor representing the intended number of rows of the
      matrix `A` as described above.

  Returns:
    A rank-2 tensor E as described above.
  """
  num_false_elem = tf.reduce_sum(
      tf.ones(shape=row_constraint.shape, dtype=tf.int32) -
      tf.cast(row_constraint, tf.int32))
  if num_false_elem == 0:
    e_matrix = tf.eye(row_dim)
  else:
    c_matrix = tf.zeros(shape=[num_false_elem, tf.size(row_constraint)])
    col_idx = tf.size(row_constraint) - num_false_elem
    updates = [1]
    for row_idx in range(num_false_elem):
      indices = [[row_idx, col_idx]]
      c_matrix = tf.tensor_scatter_nd_add(c_matrix, indices, updates)
      col_idx += 1
    e_matrix = tf.eye(row_dim) - tf.linalg.pinv(c_matrix) @ c_matrix
  return e_matrix
//  [python]
def solve_directly_for_optimal_full_batch_w(h_matrix: tf.Tensor,
                                            s_matrix: tf.Tensor) -> tf.Tensor:
  r"""Solves for W of minimal Frobenius norm satisfying WH = S.

  The resulting W matrix may be dense; this implementation leverages the
  goemetric properties of the Moore-Penrose pseudoinverse, along with the
  symmetry of the full-batch problem, to reduce complexity in computing the W
  of minimal Frobenius norm. In particular, for a vector x satisfying xA = B,
  the l^2 norm of x is at least the l^2 norm of BA^\dagger; IE, the
  Moore-Penrose pseudoinverse yields the *minimal* solution of the problem
  xA = B in l^2 norm. See Ch 2 of Campbell and Meyer's
  _Generalized Inverses of Linear Transformations_ for further discussion.

  Full-batch factorization can leverage this observation directly by applying
  it to each of the rows of W, interpreted as vectors.

  Args:
    h_matrix: rank-2 tensor representing the term H in the factorization WH = S.
    s_matrix: rank-2 tensor representing the term S in the factorization WH = S.

  Returns:
    The matrix W above; the unique rank-2 tensor of minimal Frobenius norm
    satisfying WH = S.
  """
  return tf.reshape(tf.compat.v1.to_float(s_matrix, h_matrix.dtype),
                    [-1, h_matrix.shape[1]]) @ tf.linalg.pinv(h_matrix)
//  [python]
def solve_directly_for_optimal_full_batch_w(h_matrix: tf.Tensor,
                                            s_matrix: tf.Tensor) -> tf.Tensor:
  r"""Solves for W of minimal Frobenius norm satisfying WH = S.

  The resulting W matrix may be dense; this implementation leverages the
  goemetric properties of the Moore-Penrose pseudoinverse, along with the
  symmetry of the full-batch problem, to reduce complexity in computing the W
  of minimal Frobenius norm. In particular, for a vector x satisfying xA = B,
  the l^2 norm of x is at least the l^2 norm of BA^\dagger; IE, the
  Moore-Penrose pseudoinverse yields the *minimal* solution of the problem
  xA = B in l^2 norm. See Ch 2 of Campbell and Meyer's
  _Generalized Inverses of Linear Transformations_ for further discussion.

  Full-batch factorization can leverage this observation directly by applying
  it to each of the rows of W, interpreted as vectors.

  Args:
    h_matrix: rank-2 tensor representing the term H in the factorization WH = S.
    s_matrix: rank-2 tensor representing the term S in the factorization WH = S.

  Returns:
    The matrix W above; the unique rank-2 tensor of minimal Frobenius norm
    satisfying WH = S.
  """
  return tf.reshape(tf.cast(s_matrix, h_matrix.dtype),
                    [-1, h_matrix.shape[1]]) @ tf.linalg.pinv(h_matrix)
//  [python]
def call(self, inputs):
        inputs = [t.cast(x, dtype="int32") for x in inputs]
        input_ids, input_mask, segment_ids = inputs
        bert_inputs = dict(
            input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids
        )
        if self.pooling == "first":
            pooled = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
                "pooled_output"
            ]
        elif self.pooling == "mean":
            result = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
                "sequence_output"
            ]

            mul_mask = lambda x, m: x * t.expand_dims(m, axis=-1)
            masked_reduce_mean = lambda x, m: t.reduce_sum(mul_mask(x, m), axis=1) / (
                    t.reduce_sum(m, axis=1, keepdims=True) + 1e-10)
            input_mask = t.cast(input_mask, t.float32)
            pooled = masked_reduce_mean(result, input_mask)
        else:
            raise NameError(f"Undefined pooling type (must be either first or mean, but is {self.pooling}")

        return pooled
//  [python]
def call(self, inputs):
        inputs = [K.cast(x, dtype="int32") for x in inputs]
        input_ids, input_mask, segment_ids = inputs
        bert_inputs = dict(
            input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids
        )
        if self.pooling == "first":
            pooled = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
                "pooled_output"
            ]
        elif self.pooling == "mean":
            result = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
                "sequence_output"
            ]

            mul_mask = lambda x, m: x * tf.expand_dims(m, axis=-1)
            masked_reduce_mean = lambda x, m: tf.reduce_sum(mul_mask(x, m), axis=1) / (
                    tf.reduce_sum(m, axis=1, keepdims=True) + 1e-10)
            input_mask = tf.cast(input_mask, tf.float32)
            pooled = masked_reduce_mean(result, input_mask)
        else:
            raise NameError(f"Undefined pooling type (must be either first or mean, but is {self.pooling}")

        return pooled
//  [python]
def call(self, x, mask=None, **kwargs):

        if x.dtype != tf.string:
            zero = tf.as_string(tf.zeros([1], dtype=x.dtype))
            x = tf.as_string(x, )
        else:
            zero = tf.as_string(tf.zeros([1], dtype='int32'))

        if self.vocabulary_path:
            hash_x = self.hash_table.lookup(x)
            return hash_x

        num_buckets = self.num_buckets if not self.mask_zero else self.num_buckets - 1
        try:
            hash_x = tf.string_to_hash_bucket_fast(x, num_buckets,
                                                   name=None)  # weak hash
        except AttributeError:
            hash_x = tf.strings.to_hash_bucket_fast(x, num_buckets,
                                                    name=None)  # weak hash
        if self.mask_zero:
            mask = tf.to_int64(tf.not_equal(x, zero))
            hash_x = (hash_x + 1) * mask

        return hash_x
//  [python]
def call(self, x, mask=None, **kwargs):

        if x.dtype != tf.string:
            zero = tf.as_string(tf.zeros([1], dtype=x.dtype))
            x = tf.as_string(x, )
        else:
            zero = tf.as_string(tf.zeros([1], dtype='int32'))

        if self.vocabulary_path:
            hash_x = self.hash_table.lookup(x)
            return hash_x

        num_buckets = self.num_buckets if not self.mask_zero else self.num_buckets - 1
        try:
            hash_x = tf.string_to_hash_bucket_fast(x, num_buckets,
                                                   name=None)  # weak hash
        except AttributeError:
            hash_x = tf.strings.to_hash_bucket_fast(x, num_buckets,
                                                    name=None)  # weak hash
        if self.mask_zero:
            mask = tf.cast(tf.not_equal(x, zero), dtype='int64')
            hash_x = (hash_x + 1) * mask

        return hash_x
//  [python]
def _minibatch_subsample_fn(self, inputs):
    """Randomly samples anchors for one image.

    Args:
      inputs: a list of 2 inputs. First one is a tensor of shape [num_anchors,
        num_classes] indicating targets assigned to each anchor. Second one
        is a tensor of shape [num_anchors] indicating the class weight of each
        anchor.

    Returns:
      batch_sampled_indicator: bool tensor of shape [num_anchors] indicating
        whether the anchor should be selected for loss computation.
    """
    cls_targets, cls_weights = inputs
    if self._add_background_class:
      # Set background_class bits to 0 so that the positives_indicator
      # computation would not consider background class.
      background_class = tf.zeros_like(tf.slice(cls_targets, [0, 0], [-1, 1]))
      regular_class = tf.slice(cls_targets, [0, 1], [-1, -1])
      cls_targets = tf.concat([background_class, regular_class], 1)
    positives_indicator = tf.reduce_sum(cls_targets, axis=1)
    return self._random_example_sampler.subsample(
        t.cast(cls_weights, tf.bool),
        batch_size=None,
        labels=t.cast(positives_indicator, tf.bool))