jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def extract_ho_index_by_cls_index(self, cls_index, v_gt_class_HO):
        tmp_ho_index = tf.reduce_sum(tf.multiply(v_gt_class_HO,
                                                 tf.expand_dims(tf.to_float(cls_index), axis=0)), axis=-1)
        tmp_ho_index = tf.where(tmp_ho_index > 0.0)
        tmp_ho_index = tf.squeeze(tmp_ho_index)
        return tmp_ho_index
//  [python]
def extract_ho_index_by_cls_index(self, cls_index, v_gt_class_HO):
        tmp_ho_index = tf.reduce_sum(tf.multiply(v_gt_class_HO,
                                                 tf.expand_dims(tf.cast(cls_index, tf.float32), axis=0)), axis=-1)
        tmp_ho_index = tf.where(tmp_ho_index > 0.0)
        tmp_ho_index = tf.squeeze(tmp_ho_index)
        return tmp_ho_index
//  [python]
def ldexp(
    x1: Union[tf.Tensor, tf.Variable],
    x2: Union[tf.Tensor, tf.Variable, int],
    /,
    *,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    out_dtype = x1.dtype
    x1, x2 = ivy.promote_types_of_inputs(x1, x2)
    if tf.math.reduce_any(tf.math.less(x2, 0)):
        pos_exp = t.cast(tf.math.greater_equal(x2, 0), x2.dtype) * x2
        neg_exp = t.cast(tf.math.less(x2, 0), x2.dtype) * x2
        ret = t.cast(x1, pos_exp.dtype) * tf.math.pow(2, pos_exp)
        ret = t.cast(ret, neg_exp.dtype) / tf.math.pow(2, -neg_exp)
    else:
        x2 = t.cast(x2, x1.dtype)
        ret = x1 * tf.math.pow(2, x2)
    return t.cast(ret, out_dtype)
//  [python]
def ldexp(
    x1: Union[tf.Tensor, tf.Variable],
    x2: Union[tf.Tensor, tf.Variable, int],
    /,
    *,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    out_dtype = x1.dtype
    x1, x2 = ivy.promote_types_of_inputs(x1, x2)
    if tf.math.reduce_any(tf.math.less(x2, 0)):
        pos_exp = tf.cast(tf.math.greater_equal(x2, 0), x2.dtype) * x2
        neg_exp = tf.cast(tf.math.less(x2, 0), x2.dtype) * x2
        ret = tf.cast(x1, pos_exp.dtype) * tf.math.pow(2, pos_exp)
        ret = tf.cast(ret, neg_exp.dtype) / tf.math.pow(2, -neg_exp)
    else:
        x2 = tf.cast(x2, x1.dtype)
        ret = x1 * tf.math.pow(2, x2)
    return tf.cast(ret, out_dtype)
//  [python]
def frexp(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    out: Optional[
        Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]
    ] = None,
) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]:
    e = tf.math.floor(tf.math.log(tf.math.abs(x)) / t.cast(tf.math.log(2.0), x.dtype))
    e = t.cast(e, x.dtype)
    while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1):
        e += t.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype)
    m = x / tf.math.pow(2, e)
    e = t.cast(e, tf.int32)
    return m, e
//  [python]
def frexp(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    out: Optional[
        Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]
    ] = None,
) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]:
    e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.0), x.dtype))
    e = tf.cast(e, x.dtype)
    while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1):
        e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype)
    m = x / tf.math.pow(2, e)
    e = tf.cast(e, tf.int32)
    return m, e
//  [python]
def get_wemb_mask(self, gt_obj_class, gt_obj_class_new, neighbor_num = 5):
        from networks.tools import get_neighborhood_matrix
        matrix = get_neighborhood_matrix(neighbor_num)
        matrix = tf.constant(matrix)
        available_obj_label = tf.matmul(gt_obj_class_new, matrix)
        conds = tf.reduce_sum(tf.to_int32(tf.logical_and(tf.to_bool(available_obj_label),
                                                     tf.to_bool(gt_obj_class))), axis=-1) > 0
        return conds
//  [python]
def get_wemb_mask(self, gt_obj_class, gt_obj_class_new, neighbor_num = 5):
        from networks.tools import get_neighborhood_matrix
        matrix = get_neighborhood_matrix(neighbor_num)
        matrix = tf.constant(matrix)
        available_obj_label = tf.matmul(gt_obj_class_new, matrix)
        conds = tf.reduce_sum(tf.cast(tf.logical_and(tf.cast(available_obj_label, tf.bool),
                                                     tf.cast(gt_obj_class, tf.bool)), tf.int32), axis=-1) > 0
        return conds
//  [python]
def mask_attention_for_ltr(attention_scores, attention_mask):
    """
    Mask attention so that we're only predicting going forward
    :param attention_scores: [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
    :param attention_mask [query_length, key_length]
    :return: masked attention
    """
    # 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.
    mask = attention_mask[None, None]
    return attention_scores * mask - t.cast(1e10, attention_scores.dtype) * (1 - mask)
//  [python]
def mask_attention_for_ltr(attention_scores, attention_mask):
    """
    Mask attention so that we're only predicting going forward
    :param attention_scores: [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
    :param attention_mask [query_length, key_length]
    :return: masked attention
    """
    # 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.
    mask = attention_mask[None, None]
    return attention_scores * mask - tf.cast(1e10, attention_scores.dtype) * (1 - mask)
//  [python]
def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
    """
    Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is
    True` else after `sep_token_id`.
    """

    assert shape_list(sep_token_indices)[1] == 2, "`input_ids` should have two dimensions"
    question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1][:, None]
    # bool attention mask with True in locations of global attention
    attention_mask = tf.expand_dims(tf.range(input_ids_shape[1]), axis=0)
    attention_mask = tf.tile(attention_mask, (input_ids_shape[0], 1))
    if before_sep_token is True:
        question_end_index = tf.tile(question_end_index, (1, input_ids_shape[1]))
        attention_mask = tf.to_int32(attention_mask < question_end_index)
    else:
        # last token is separation token and should not be counted and in the middle are two separation tokens
        question_end_index = tf.tile(question_end_index + 1, (1, input_ids_shape[1]))
        attention_mask = tf.to_int32(
            attention_mask > question_end_index,
        ) * tf.to_int32(attention_mask < input_ids_shape[-1])

    return attention_mask
//  [python]
def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
    """
    Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is
    True` else after `sep_token_id`.
    """

    assert shape_list(sep_token_indices)[1] == 2, "`input_ids` should have two dimensions"
    question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1][:, None]
    # bool attention mask with True in locations of global attention
    attention_mask = tf.expand_dims(tf.range(input_ids_shape[1]), axis=0)
    attention_mask = tf.tile(attention_mask, (input_ids_shape[0], 1))
    if before_sep_token is True:
        question_end_index = tf.tile(question_end_index, (1, input_ids_shape[1]))
        attention_mask = tf.cast(attention_mask < question_end_index, dtype=question_end_index.dtype)
    else:
        # last token is separation token and should not be counted and in the middle are two separation tokens
        question_end_index = tf.tile(question_end_index + 1, (1, input_ids_shape[1]))
        attention_mask = tf.cast(
            attention_mask > question_end_index,
            dtype=question_end_index.dtype,
        ) * tf.cast(attention_mask < input_ids_shape[-1], dtype=question_end_index.dtype)

    return attention_mask
//  [python]
def _auxiliary_loss(self,
                        h_states,
                        click_seq,
                        noclick_seq,
                        mask,
                        dtype=tf.float32,
                        stag=None):
        mask = tf.cast(mask, dtype=dtype)
        click_input_ = tf.concat([h_states, click_seq], -1)
        noclick_input_ = tf.concat([h_states, noclick_seq], -1)
        with tf.variable_scope('auxiliary_net'):
            click_prop_ = self._auxiliary_net(click_input_, stag=stag)[:, :, 0]
            noclick_prop_ = self._auxiliary_net(noclick_input_,
                                                stag=stag)[:, :, 0]

        click_loss_ = -tf.reshape(tf.log(click_prop_),
                                  [-1, tf.shape(click_seq)[1]]) * mask
        noclick_loss_ = -tf.reshape(tf.log(1.0 - noclick_prop_),
                                    [-1, tf.shape(noclick_seq)[1]]) * mask
        loss_ = tf.reduce_mean(click_loss_ + noclick_loss_)
        return loss_
//  [python]
def _mask_id_xent_loss(loss_dict: Dict[str, Any], labels: Dict[str, Any],
                       outputs: Dict[str, Any]):
  """Mask ID loss.

  This method adds the mask ID loss term to loss_dict directly.

  Args:
    loss_dict: A dictionary for the loss. The values are loss scalars.
    labels: The label dictionary.
    outputs: The output dictionary.
  """
  # (B, N, H, W)
  mask_gt = labels["masks"]
  # B, H, W, N
  mask_id_logits = outputs["instance_output"]["mask_id_logits"]
  # B, N, N
  matched_matrix = outputs["instance_output"]["matched_mask"]
  # B, N
  gt_to_pred_id = t.cast(t.math.argmax(matched_matrix, axis=1), t.float32)
  # B, H, W
  mask_id_labels = t.cast(
      t.einsum("bnhw,bn->bhw", mask_gt, gt_to_pred_id), t.int32)
  loss_map = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=mask_id_labels, logits=mask_id_logits)
  valid_mask = tf.reduce_sum(mask_gt, axis=1)
  loss_mask_id = (
      (tf.reduce_sum(loss_map * valid_mask, axis=[1, 2]) + EPSILON) /
      (tf.reduce_sum(valid_mask, axis=[1, 2]) + EPSILON))
  loss_dict["loss_mask_id"] = tf.reduce_mean(loss_mask_id)
//  [python]
def _mask_id_xent_loss(loss_dict: Dict[str, Any], labels: Dict[str, Any],
                       outputs: Dict[str, Any]):
  """Mask ID loss.

  This method adds the mask ID loss term to loss_dict directly.

  Args:
    loss_dict: A dictionary for the loss. The values are loss scalars.
    labels: The label dictionary.
    outputs: The output dictionary.
  """
  # (B, N, H, W)
  mask_gt = labels["masks"]
  # B, H, W, N
  mask_id_logits = outputs["instance_output"]["mask_id_logits"]
  # B, N, N
  matched_matrix = outputs["instance_output"]["matched_mask"]
  # B, N
  gt_to_pred_id = tf.cast(tf.math.argmax(matched_matrix, axis=1), tf.float32)
  # B, H, W
  mask_id_labels = tf.cast(
      tf.einsum("bnhw,bn->bhw", mask_gt, gt_to_pred_id), tf.int32)
  loss_map = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=mask_id_labels, logits=mask_id_logits)
  valid_mask = tf.reduce_sum(mask_gt, axis=1)
  loss_mask_id = (
      (tf.reduce_sum(loss_map * valid_mask, axis=[1, 2]) + EPSILON) /
      (tf.reduce_sum(valid_mask, axis=[1, 2]) + EPSILON))
  loss_dict["loss_mask_id"] = tf.reduce_mean(loss_mask_id)
//  [python]
def pooled_output(self, clf_token):
        """
        Extract pooled output given a token that says where we should look
        :param clf_token:
        :return:
        """
        pool_idx = t.cast(t.argmax(t.cast(t.equal(self.input_ids, clf_token), t.float32), 1), t.int32)
        return t.gather(self.hidden_state, t.range(self.batch_size, dtype=t.int32) * self.seq_length + pool_idx)
//  [python]
def pooled_output(self, clf_token):
        """
        Extract pooled output given a token that says where we should look
        :param clf_token:
        :return:
        """
        pool_idx = tf.cast(tf.argmax(tf.cast(tf.equal(self.input_ids, clf_token), tf.float32), 1), tf.int32)
        return tf.gather(self.hidden_state, tf.range(self.batch_size, dtype=tf.int32) * self.seq_length + pool_idx)
//  [python]
def _postprocess(self, outputs: typing.NestedTensorDict):
    """Post-process (filtering) the outputs.

    Args:
      outputs: A dictionary of outputs.
    These following fields are added to outputs["postprocessed"]:
      "classes": A (B,N) integer tensor for the class ids.
      "binary_masks": A (B, H, W, N) tensor for the N binarized 0/1 masks. Masks
        for void cls are set to zero.
      "confidence": A (B, N) float tensor for the confidence of "classes".
      "mask_area": A (B, N) float tensor for the area of each mask. They are
        used in inference / visualization.
    """
    # Get postprocessed outputs
    outputs["postprocessed"] = {}

    ## Masks:
    mask_id_prob = outputs["instance_output"]["mask_id_prob"]
    mask_max_prob = tf.reduce_max(mask_id_prob, axis=-1, keepdims=True)
    thresholded_binary_masks = t.cast(
        tf.math.logical_and(
            tf.equal(mask_max_prob, mask_id_prob),
            tf.greater_equal(mask_max_prob, self._mask_threshold)), tf.float32)
    area = tf.reduce_sum(thresholded_binary_masks, axis=(1, 2))  # (B, N)
    ## Classification:
    cls_prob = outputs["instance_output"]["cls_prob"]
    cls_max_prob = tf.reduce_max(cls_prob, axis=-1)  # B, N
    cls_max_id = t.cast(tf.argmax(cls_prob, axis=-1), tf.float32)  # B, N

    ## filtering
    c = utilities.resolve_shape(cls_prob)[2]
    non_void = tf.reduce_all(
        tf.stack(
            [
                tf.greater_equal(area, self._filter_area),  # mask large enough.
                tf.not_equal(cls_max_id, 0),  # class-0 is for non-object.
                tf.not_equal(cls_max_id,
                             c - 1),  # class-(c-1) is for background (last).
                tf.greater_equal(cls_max_prob,
                                 self._class_threshold)  # prob >= thr
            ],
            axis=-1),
        axis=-1)
    non_void = t.cast(non_void, tf.float32)

    # Storing
    outputs["postprocessed"]["classes"] = t.cast(cls_max_id * non_void,
                                                  tf.int32)
    b, n = utilities.resolve_shape(non_void)
    outputs["postprocessed"]["binary_masks"] = (
        thresholded_binary_masks * t.reshape(non_void, (b, 1, 1, n)))
    outputs["postprocessed"]["confidence"] = cls_max_prob
    outputs["postprocessed"]["mask_area"] = area
//  [python]
def _postprocess(self, outputs: typing.NestedTensorDict):
    """Post-process (filtering) the outputs.

    Args:
      outputs: A dictionary of outputs.
    These following fields are added to outputs["postprocessed"]:
      "classes": A (B,N) integer tensor for the class ids.
      "binary_masks": A (B, H, W, N) tensor for the N binarized 0/1 masks. Masks
        for void cls are set to zero.
      "confidence": A (B, N) float tensor for the confidence of "classes".
      "mask_area": A (B, N) float tensor for the area of each mask. They are
        used in inference / visualization.
    """
    # Get postprocessed outputs
    outputs["postprocessed"] = {}

    ## Masks:
    mask_id_prob = outputs["instance_output"]["mask_id_prob"]
    mask_max_prob = tf.reduce_max(mask_id_prob, axis=-1, keepdims=True)
    thresholded_binary_masks = tf.cast(
        tf.math.logical_and(
            tf.equal(mask_max_prob, mask_id_prob),
            tf.greater_equal(mask_max_prob, self._mask_threshold)), tf.float32)
    area = tf.reduce_sum(thresholded_binary_masks, axis=(1, 2))  # (B, N)
    ## Classification:
    cls_prob = outputs["instance_output"]["cls_prob"]
    cls_max_prob = tf.reduce_max(cls_prob, axis=-1)  # B, N
    cls_max_id = tf.cast(tf.argmax(cls_prob, axis=-1), tf.float32)  # B, N

    ## filtering
    c = utilities.resolve_shape(cls_prob)[2]
    non_void = tf.reduce_all(
        tf.stack(
            [
                tf.greater_equal(area, self._filter_area),  # mask large enough.
                tf.not_equal(cls_max_id, 0),  # class-0 is for non-object.
                tf.not_equal(cls_max_id,
                             c - 1),  # class-(c-1) is for background (last).
                tf.greater_equal(cls_max_prob,
                                 self._class_threshold)  # prob >= thr
            ],
            axis=-1),
        axis=-1)
    non_void = tf.cast(non_void, tf.float32)

    # Storing
    outputs["postprocessed"]["classes"] = tf.cast(cls_max_id * non_void,
                                                  tf.int32)
    b, n = utilities.resolve_shape(non_void)
    outputs["postprocessed"]["binary_masks"] = (
        thresholded_binary_masks * tf.reshape(non_void, (b, 1, 1, n)))
    outputs["postprocessed"]["confidence"] = cls_max_prob
    outputs["postprocessed"]["mask_area"] = area
//  [python]
def lm_loss(self):
        """
        :return: stuff
        """
        target_ids_flat = tf.reshape(self.target_ids, [-1])

        # 1 if it's valid and 0 otherwise.
        label_weights = tf.to_float(tf.not_equal(target_ids_flat, self.pad_token_id))

        # [batch_size * seq_length, vocab_size]
        one_hot_labels = tf.one_hot(target_ids_flat,
                                    depth=self.config.vocab_size,
                                    dtype=self.logits_flat.dtype)

        # [batch_size * seq_length, vocab_size]
        logprobs_flat = tf.nn.log_softmax(self.logits_flat, axis=-1)

        per_example_loss = -tf.reduce_sum(logprobs_flat * one_hot_labels, axis=[-1])

        # per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits_flat, labels=target_ids_flat)

        numerator = tf.reduce_sum(label_weights * per_example_loss)
        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = numerator / denominator
        return loss
//  [python]
def lm_loss(self):
        """
        :return: stuff
        """
        target_ids_flat = tf.reshape(self.target_ids, [-1])

        # 1 if it's valid and 0 otherwise.
        label_weights = tf.cast(tf.not_equal(target_ids_flat, self.pad_token_id), dtype=self.logits_flat.dtype)

        # [batch_size * seq_length, vocab_size]
        one_hot_labels = tf.one_hot(target_ids_flat,
                                    depth=self.config.vocab_size,
                                    dtype=self.logits_flat.dtype)

        # [batch_size * seq_length, vocab_size]
        logprobs_flat = tf.nn.log_softmax(self.logits_flat, axis=-1)

        per_example_loss = -tf.reduce_sum(logprobs_flat * one_hot_labels, axis=[-1])

        # per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits_flat, labels=target_ids_flat)

        numerator = tf.reduce_sum(label_weights * per_example_loss)
        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = numerator / denominator
        return loss
//  [python]
def _instance_discrimination_loss(loss_dict: Dict[str, Any],
                                  labels: Dict[str, Any],
                                  outputs: Dict[str, Any],
                                  tau: float = gin.REQUIRED):
  """Instance discrimination loss.

  This method adds the ID loss term to loss_dict directly.

  Args:
    loss_dict: A dictionary for the loss. The values are loss scalars.
    labels: The label dictionary.
    outputs: The output dictionary.
    tau: The temperature term in the loss
  """
  # The normalized feature, shape=(B, H/4, W/4, D)
  g = outputs["max_deep_lab"]["pixel_space_normalized_feature"]
  b, h, w = utilities.resolve_shape(g)[:3]
  # The ground-truth masks, shape=(B, N, H, W) --> (B, N, H/4, W/4)
  m = labels["masks"]
  m = tf.image.resize(
      tf.transpose(m, (0, 2, 3, 1)), (h, w),
      tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  m = tf.transpose(m, (0, 3, 1, 2))
  # The number of ground-truth instance (K), shape=(B,)
  num = labels["num_instance"]
  n = utilities.resolve_shape(m)[1]  # max number of predictions
  # is_void[b, i] = 1 if instance i in batch b is a padded slot.
  is_void = tf.to_float(tf.expand_dims(tf.range(n), 0))  # (1, n)
  is_void = tf.to_float(
      tf.math.greater_equal(is_void, tf.expand_dims(num, 1)))

  # (B, N, D)
  t = tf.math.l2_normalize(tf.einsum("bhwd,bnhw->bnd", g, m), axis=-1)
  inst_dist_logits = tf.einsum("bhwd,bid->bhwi", g, t) / tau  # (B, H, W, N)
  inst_dist_logits = inst_dist_logits - 100. * tf.reshape(is_void, (b, 1, 1, n))
  mask_id = tf.to_int32(
      tf.einsum("bnhw,n->bhw", m, tf.range(n, dtype=tf.float32)))
  loss_map = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=mask_id, logits=inst_dist_logits)  # B, H, W
  valid_mask = tf.reduce_sum(m, axis=1)
  loss_inst_dist = (
      (tf.reduce_sum(loss_map * valid_mask, axis=[1, 2]) + EPSILON) /
      (tf.reduce_sum(valid_mask, axis=[1, 2]) + EPSILON))
  loss_dict["loss_inst_dist"] = tf.reduce_mean(loss_inst_dist)
//  [python]
def _instance_discrimination_loss(loss_dict: Dict[str, Any],
                                  labels: Dict[str, Any],
                                  outputs: Dict[str, Any],
                                  tau: float = gin.REQUIRED):
  """Instance discrimination loss.

  This method adds the ID loss term to loss_dict directly.

  Args:
    loss_dict: A dictionary for the loss. The values are loss scalars.
    labels: The label dictionary.
    outputs: The output dictionary.
    tau: The temperature term in the loss
  """
  # The normalized feature, shape=(B, H/4, W/4, D)
  g = outputs["max_deep_lab"]["pixel_space_normalized_feature"]
  b, h, w = utilities.resolve_shape(g)[:3]
  # The ground-truth masks, shape=(B, N, H, W) --> (B, N, H/4, W/4)
  m = labels["masks"]
  m = tf.image.resize(
      tf.transpose(m, (0, 2, 3, 1)), (h, w),
      tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  m = tf.transpose(m, (0, 3, 1, 2))
  # The number of ground-truth instance (K), shape=(B,)
  num = labels["num_instance"]
  n = utilities.resolve_shape(m)[1]  # max number of predictions
  # is_void[b, i] = 1 if instance i in batch b is a padded slot.
  is_void = tf.cast(tf.expand_dims(tf.range(n), 0), tf.float32)  # (1, n)
  is_void = tf.cast(
      tf.math.greater_equal(is_void, tf.expand_dims(num, 1)), tf.float32)

  # (B, N, D)
  t = tf.math.l2_normalize(tf.einsum("bhwd,bnhw->bnd", g, m), axis=-1)
  inst_dist_logits = tf.einsum("bhwd,bid->bhwi", g, t) / tau  # (B, H, W, N)
  inst_dist_logits = inst_dist_logits - 100. * tf.reshape(is_void, (b, 1, 1, n))
  mask_id = tf.cast(
      tf.einsum("bnhw,n->bhw", m, tf.range(n, dtype=tf.float32)), tf.int32)
  loss_map = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=mask_id, logits=inst_dist_logits)  # B, H, W
  valid_mask = tf.reduce_sum(m, axis=1)
  loss_inst_dist = (
      (tf.reduce_sum(loss_map * valid_mask, axis=[1, 2]) + EPSILON) /
      (tf.reduce_sum(valid_mask, axis=[1, 2]) + EPSILON))
  loss_dict["loss_inst_dist"] = tf.reduce_mean(loss_inst_dist)
//  [python]
def get_masks(self, requests_state):
        request_mask_s = t.cast(t.reduce_sum(requests_state, axis=2) > 0, t.float32)
        request_mask_l = t.repeat(request_mask_s, repeats=self.n_veh, axis=1)
        
        request_masks = {"s": t.stop_gradient(request_mask_s),
                         "l": t.stop_gradient(request_mask_l)}
        
        return request_masks
//  [python]
def get_masks(self, requests_state):
        request_mask_s = tf.cast(tf.reduce_sum(requests_state, axis=2) > 0, tf.float32)
        request_mask_l = tf.repeat(request_mask_s, repeats=self.n_veh, axis=1)
        
        request_masks = {"s": tf.stop_gradient(request_mask_s),
                         "l": tf.stop_gradient(request_mask_l)}
        
        return request_masks
//  [python]
def get_minibatch_tf(self): # => images, labels
        with tf.name_scope('SyntheticDataset'):
            shrink = tf.to_int32(tf.pow(2.0, tf.to_float(self._tf_lod_var)))
            shape = [self.shape[0], self.shape[1] // shrink, self.shape[2] // shrink]
            images = self._generate_images(self._tf_minibatch_var, self._tf_lod_var, shape)
            labels = self._generate_labels(self._tf_minibatch_var)
            return images, labels
//  [python]
def get_minibatch_tf(self): # => images, labels
        with tf.name_scope('SyntheticDataset'):
            shrink = tf.cast(2.0 ** tf.cast(self._tf_lod_var, tf.float32), tf.int32)
            shape = [self.shape[0], self.shape[1] // shrink, self.shape[2] // shrink]
            images = self._generate_images(self._tf_minibatch_var, self._tf_lod_var, shape)
            labels = self._generate_labels(self._tf_minibatch_var)
            return images, labels
//  [python]
def make_input_fn(is_training):
    """Construct training input_fn that uses synthetic data."""

    def input_fn(params):
      """Returns dummy input batches for training."""

      # Estimator passes batch_size during training and eval_batch_size during
      # eval.
      batch_size = (
          params["batch_size"] if is_training else
          params.get("eval_batch_size") or params["batch_size"])
      num_users = params["num_users"]
      num_items = params["num_items"]

      users = tf.random.uniform([batch_size, 1],
                                dtype=tf.int32,
                                minval=0,
                                maxval=num_users)
      items = tf.random.uniform([batch_size, 1],
                                dtype=tf.int32,
                                minval=0,
                                maxval=num_items)

      if is_training:
        valid_point_mask = tf.to_int32(
            tf.random.uniform([batch_size, 1],
                              dtype=tf.int32,
                              minval=0,
                              maxval=2)) > 0
        labels = tf.to_int32(
            tf.random.uniform([batch_size, 1],
                              dtype=tf.int32,
                              minval=0,
                              maxval=2)) > 0
        data = {
            movielens.USER_COLUMN: users,
            movielens.ITEM_COLUMN: items,
            rconst.VALID_POINT_MASK: valid_point_mask,
        }, labels
      else:
        dupe_mask = tf.to_int32(
            tf.random.uniform([batch_size, 1],
                              dtype=tf.int32,
                              minval=0,
                              maxval=2)) > 0
        data = {
            movielens.USER_COLUMN: users,
            movielens.ITEM_COLUMN: items,
            rconst.DUPLICATE_MASK: dupe_mask,
        }

      dataset = tf.data.Dataset.from_tensors(data).repeat(
          rconst.SYNTHETIC_BATCHES_PER_EPOCH * params["batches_per_step"])
      dataset = dataset.prefetch(32)
      return dataset

    return input_fn
//  [python]
def make_input_fn(is_training):
    """Construct training input_fn that uses synthetic data."""

    def input_fn(params):
      """Returns dummy input batches for training."""

      # Estimator passes batch_size during training and eval_batch_size during
      # eval.
      batch_size = (
          params["batch_size"] if is_training else
          params.get("eval_batch_size") or params["batch_size"])
      num_users = params["num_users"]
      num_items = params["num_items"]

      users = tf.random.uniform([batch_size, 1],
                                dtype=tf.int32,
                                minval=0,
                                maxval=num_users)
      items = tf.random.uniform([batch_size, 1],
                                dtype=tf.int32,
                                minval=0,
                                maxval=num_items)

      if is_training:
        valid_point_mask = tf.cast(
            tf.random.uniform([batch_size, 1],
                              dtype=tf.int32,
                              minval=0,
                              maxval=2), tf.bool)
        labels = tf.cast(
            tf.random.uniform([batch_size, 1],
                              dtype=tf.int32,
                              minval=0,
                              maxval=2), tf.bool)
        data = {
            movielens.USER_COLUMN: users,
            movielens.ITEM_COLUMN: items,
            rconst.VALID_POINT_MASK: valid_point_mask,
        }, labels
      else:
        dupe_mask = tf.cast(
            tf.random.uniform([batch_size, 1],
                              dtype=tf.int32,
                              minval=0,
                              maxval=2), tf.bool)
        data = {
            movielens.USER_COLUMN: users,
            movielens.ITEM_COLUMN: items,
            rconst.DUPLICATE_MASK: dupe_mask,
        }

      dataset = tf.data.Dataset.from_tensors(data).repeat(
          rconst.SYNTHETIC_BATCHES_PER_EPOCH * params["batches_per_step"])
      dataset = dataset.prefetch(32)
      return dataset

    return input_fn
//  [python]
def accuracy(self):
        if self.label_data.shape[1]>1:
            pre_lables=tf.argmax(self.pred,axis=1)
            data_lables=tf.argmax(self.label_data,axis=1)
        else:
            pre_lables=tf.round(self.pred)
            data_lables=tf.round(self.label_data)
        return tf.reduce_mean(tf.to_float(tf.equal(pre_lables,data_lables)))
//  [python]
def accuracy(self):
        if self.label_data.shape[1]>1:
            pre_lables=tf.argmax(self.pred,axis=1)
            data_lables=tf.argmax(self.label_data,axis=1)
        else:
            pre_lables=tf.round(self.pred)
            data_lables=tf.round(self.label_data)
        return tf.reduce_mean(tf.cast(tf.equal(pre_lables,data_lables),tf.float32))
//  [python]
def deserialize(serialized_data, batch_size=None, is_training=True):
    """Convert serialized TFRecords into tensors.

    Args:
      serialized_data: A tensor containing serialized records.
      batch_size: The data arrives pre-batched, so batch size is needed to
        deserialize the data.
      is_training: Boolean, whether data to deserialize to training data or
        evaluation data.
    """

    def _get_feature_map(batch_size, is_training=True):
      """Returns data format of the serialized tf record file."""

      if is_training:
        return {
            movielens.USER_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            movielens.ITEM_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            rconst.VALID_POINT_MASK:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            "labels":
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64)
        }
      else:
        return {
            movielens.USER_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            movielens.ITEM_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            rconst.DUPLICATE_MASK:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64)
        }

    features = tf.io.parse_single_example(
        serialized_data, _get_feature_map(batch_size, is_training=is_training))
    users = t.cast(features[movielens.USER_COLUMN], rconst.USER_DTYPE)
    items = t.cast(features[movielens.ITEM_COLUMN], rconst.ITEM_DTYPE)

    if is_training:
      valid_point_mask = t.cast(features[rconst.VALID_POINT_MASK], tf.bool)
      fake_dup_mask = tf.zeros_like(users)
      return {
          movielens.USER_COLUMN:
              users,
          movielens.ITEM_COLUMN:
              items,
          rconst.VALID_POINT_MASK:
              valid_point_mask,
          rconst.TRAIN_LABEL_KEY:
              tf.reshape(t.cast(features["labels"], tf.bool), (batch_size, 1)),
          rconst.DUPLICATE_MASK:
              fake_dup_mask
      }
    else:
      labels = t.cast(tf.zeros_like(users), tf.bool)
      fake_valid_pt_mask = t.cast(tf.zeros_like(users), tf.bool)
      return {
          movielens.USER_COLUMN:
              users,
          movielens.ITEM_COLUMN:
              items,
          rconst.DUPLICATE_MASK:
              t.cast(features[rconst.DUPLICATE_MASK], tf.bool),
          rconst.VALID_POINT_MASK:
              fake_valid_pt_mask,
          rconst.TRAIN_LABEL_KEY:
              labels
      }
//  [python]
def deserialize(serialized_data, batch_size=None, is_training=True):
    """Convert serialized TFRecords into tensors.

    Args:
      serialized_data: A tensor containing serialized records.
      batch_size: The data arrives pre-batched, so batch size is needed to
        deserialize the data.
      is_training: Boolean, whether data to deserialize to training data or
        evaluation data.
    """

    def _get_feature_map(batch_size, is_training=True):
      """Returns data format of the serialized tf record file."""

      if is_training:
        return {
            movielens.USER_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            movielens.ITEM_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            rconst.VALID_POINT_MASK:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            "labels":
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64)
        }
      else:
        return {
            movielens.USER_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            movielens.ITEM_COLUMN:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64),
            rconst.DUPLICATE_MASK:
                tf.io.FixedLenFeature([batch_size, 1], dtype=tf.int64)
        }

    features = tf.io.parse_single_example(
        serialized_data, _get_feature_map(batch_size, is_training=is_training))
    users = tf.cast(features[movielens.USER_COLUMN], rconst.USER_DTYPE)
    items = tf.cast(features[movielens.ITEM_COLUMN], rconst.ITEM_DTYPE)

    if is_training:
      valid_point_mask = tf.cast(features[rconst.VALID_POINT_MASK], tf.bool)
      fake_dup_mask = tf.zeros_like(users)
      return {
          movielens.USER_COLUMN:
              users,
          movielens.ITEM_COLUMN:
              items,
          rconst.VALID_POINT_MASK:
              valid_point_mask,
          rconst.TRAIN_LABEL_KEY:
              tf.reshape(tf.cast(features["labels"], tf.bool), (batch_size, 1)),
          rconst.DUPLICATE_MASK:
              fake_dup_mask
      }
    else:
      labels = tf.cast(tf.zeros_like(users), tf.bool)
      fake_valid_pt_mask = tf.cast(tf.zeros_like(users), tf.bool)
      return {
          movielens.USER_COLUMN:
              users,
          movielens.ITEM_COLUMN:
              items,
          rconst.DUPLICATE_MASK:
              tf.cast(features[rconst.DUPLICATE_MASK], tf.bool),
          rconst.VALID_POINT_MASK:
              fake_valid_pt_mask,
          rconst.TRAIN_LABEL_KEY:
              labels
      }
//  [python]
def call(
      self,
      inputs: types.NestedTensor,
      outer_rank: int,
      training: bool = False,
      mask: Optional[types.NestedTensor] = None,
  ) -> types.NestedTensor:
    if inputs.dtype != self._sample_spec.dtype:  # pytype: disable=attribute-error
      raise ValueError(
          'Inputs to TanhNormalProjectionNetwork must match the '
          'sample_spec.dtype.'
      )

    if mask is not None:
      raise NotImplementedError(
          'TanhNormalProjectionNetwork does not yet implement action masking; '
          'got mask={}'.format(mask)
      )

    # outer_rank is needed because the projection is not done on the raw
    # observations so getting the outer rank is hard as there is no spec to
    # compare to.
    batch_squash = network_utils.BatchSquash(outer_rank)
    inputs = batch_squash.flatten(inputs)

    means_and_stds = self._projection_layer(inputs, training=training)
    means, stds = tf.split(means_and_stds, num_or_size_splits=2, axis=-1)
    means = tf.reshape(means, [-1] + self._sample_spec.shape.as_list())
    means = t.cast(means, self._sample_spec.dtype)

    if self._std_transform is not None:
      stds = self._std_transform(stds)
    stds = t.cast(stds, self._sample_spec.dtype)

    means = batch_squash.unflatten(means)
    stds = batch_squash.unflatten(stds)

    return self.output_spec.build_distribution(loc=means, scale_diag=stds), ()
//  [python]
def call(
      self,
      inputs: types.NestedTensor,
      outer_rank: int,
      training: bool = False,
      mask: Optional[types.NestedTensor] = None,
  ) -> types.NestedTensor:
    if inputs.dtype != self._sample_spec.dtype:  # pytype: disable=attribute-error
      raise ValueError(
          'Inputs to TanhNormalProjectionNetwork must match the '
          'sample_spec.dtype.'
      )

    if mask is not None:
      raise NotImplementedError(
          'TanhNormalProjectionNetwork does not yet implement action masking; '
          'got mask={}'.format(mask)
      )

    # outer_rank is needed because the projection is not done on the raw
    # observations so getting the outer rank is hard as there is no spec to
    # compare to.
    batch_squash = network_utils.BatchSquash(outer_rank)
    inputs = batch_squash.flatten(inputs)

    means_and_stds = self._projection_layer(inputs, training=training)
    means, stds = tf.split(means_and_stds, num_or_size_splits=2, axis=-1)
    means = tf.reshape(means, [-1] + self._sample_spec.shape.as_list())
    means = tf.cast(means, self._sample_spec.dtype)

    if self._std_transform is not None:
      stds = self._std_transform(stds)
    stds = tf.cast(stds, self._sample_spec.dtype)

    means = batch_squash.unflatten(means)
    stds = batch_squash.unflatten(stds)

    return self.output_spec.build_distribution(loc=means, scale_diag=stds), ()
//  [python]
def _maybe_apply_per_example_weight(
      self,
      observation: tf.Tensor,
      reward: tf.Tensor,
      weights: Optional[tf.Tensor],
  ) -> Tuple[tf.Tensor, tf.Tensor]:
    """Optionally applies per-example weight to observation and rewards."""
    if weights is None:
      return (observation, reward)
    else:
      w_sqrt = tf.sqrt(tf.cast(weights, dtype=self._dtype))
      return (tf.reshape(w_sqrt, [-1, 1]) * observation, w_sqrt * reward)
//  [python]
def _generate_mask_and_prune_ops(self, v):
        if self._retain_sparsity:
            th = self._eps
        else:
            th = -1.0
        
        self._mask[v] = tf.to_float(tf.abs(v) > th)

        if isinstance(v, tf.Tensor):
            return
        
        if self._retain_sparsity:
            prune_weight_op = tf.assign(v, v * self._mask[v])
            self._prune_weight_ops.append(prune_weight_op)
//  [python]
def _generate_mask_and_prune_ops(self, v):
        if self._retain_sparsity:
            th = self._eps
        else:
            th = -1.0
        
        self._mask[v] = tf.cast(tf.abs(v) > th, tf.float32)

        if isinstance(v, tf.Tensor):
            return
        
        if self._retain_sparsity:
            prune_weight_op = tf.assign(v, v * self._mask[v])
            self._prune_weight_ops.append(prune_weight_op)
//  [python]
def call(self, inputs):
        inputs, unbatched, _ = convert_to_ragged_batch(inputs)

        skip_masks = None
        if self.skip_list:
            skip_masks = self.StaticHashTable.lookup(inputs.flat_values)
        elif self.skip_fn:
            skip_masks = tf.map_fn(
                self.skip_fn, inputs.flat_values, fn_output_signature="bool"
            )
        elif self.skip_py_fn:

            def string_fn(token):
                return self.skip_py_fn(token.numpy().decode("utf-8"))

            def int_fn(token):
                return self.skip_py_fn(token.numpy())

            py_fn = string_fn if inputs.dtype == tf.string else int_fn

            skip_masks = tf.map_fn(
                lambda x: tf.py_function(py_fn, [x], "bool"),
                inputs.flat_values,
                fn_output_signature="bool",
            )

        positions_flat = tf.range(tf.size(inputs.flat_values))
        positions = inputs.with_flat_values(positions_flat)
        if skip_masks is not None:
            skip_masks = tf.logical_not(skip_masks)
            skip_masks.set_shape([None])
            positions = tf.ragged.boolean_mask(
                positions, inputs.with_flat_values(skip_masks)
            )

        # Figure out how many we are going to select.
        token_counts = t.cast(positions.row_lengths(), "float32")
        num_to_select = tf.random.stateless_binomial(
            shape=tf.shape(token_counts),
            seed=self._generator.make_seeds()[:, 0],
            counts=token_counts,
            probs=self.rate,
        )
        if self.max_deletions is not None:
            num_to_select = tf.math.minimum(num_to_select, self.max_deletions)
        num_to_select = t.cast(num_to_select, "int64")

        # Shuffle and trim to items that are going to be selected.
        def _shuffle_and_trim(x):
            positions, top_n = x
            shuffled = tf.random.shuffle(positions, seed=self.seed)
            return shuffled[:top_n]

        selected_for_mask = tf.map_fn(
            _shuffle_and_trim,
            (positions, num_to_select),
            fn_output_signature=tf.RaggedTensorSpec(
                ragged_rank=positions.ragged_rank - 1, dtype=positions.dtype
            ),
        )
        selected_for_mask.flat_values.set_shape([None])

        # Construct the mask which is a boolean RT
        # Scatter 0's to positions that have been selector for deletion.
        update_values = tf.zeros_like(selected_for_mask.flat_values, "int32")
        update_indices = selected_for_mask.flat_values
        update_indices = tf.expand_dims(update_indices, -1)
        update_indices = t.cast(update_indices, "int32")
        mask_flat = tf.ones_like(inputs.flat_values, dtype="int32")
        mask_flat = tf.tensor_scatter_nd_update(
            mask_flat, update_indices, update_values
        )
        mask = t.cast(inputs.with_flat_values(mask_flat), "bool")

        inputs = tf.ragged.boolean_mask(inputs, mask)

        if unbatched:
            inputs = tf.squeeze(inputs, axis=0)

        return inputs
//  [python]
def call(self, inputs):
        inputs, unbatched, _ = convert_to_ragged_batch(inputs)

        skip_masks = None
        if self.skip_list:
            skip_masks = self.StaticHashTable.lookup(inputs.flat_values)
        elif self.skip_fn:
            skip_masks = tf.map_fn(
                self.skip_fn, inputs.flat_values, fn_output_signature="bool"
            )
        elif self.skip_py_fn:

            def string_fn(token):
                return self.skip_py_fn(token.numpy().decode("utf-8"))

            def int_fn(token):
                return self.skip_py_fn(token.numpy())

            py_fn = string_fn if inputs.dtype == tf.string else int_fn

            skip_masks = tf.map_fn(
                lambda x: tf.py_function(py_fn, [x], "bool"),
                inputs.flat_values,
                fn_output_signature="bool",
            )

        positions_flat = tf.range(tf.size(inputs.flat_values))
        positions = inputs.with_flat_values(positions_flat)
        if skip_masks is not None:
            skip_masks = tf.logical_not(skip_masks)
            skip_masks.set_shape([None])
            positions = tf.ragged.boolean_mask(
                positions, inputs.with_flat_values(skip_masks)
            )

        # Figure out how many we are going to select.
        token_counts = tf.cast(positions.row_lengths(), "float32")
        num_to_select = tf.random.stateless_binomial(
            shape=tf.shape(token_counts),
            seed=self._generator.make_seeds()[:, 0],
            counts=token_counts,
            probs=self.rate,
        )
        if self.max_deletions is not None:
            num_to_select = tf.math.minimum(num_to_select, self.max_deletions)
        num_to_select = tf.cast(num_to_select, "int64")

        # Shuffle and trim to items that are going to be selected.
        def _shuffle_and_trim(x):
            positions, top_n = x
            shuffled = tf.random.shuffle(positions, seed=self.seed)
            return shuffled[:top_n]

        selected_for_mask = tf.map_fn(
            _shuffle_and_trim,
            (positions, num_to_select),
            fn_output_signature=tf.RaggedTensorSpec(
                ragged_rank=positions.ragged_rank - 1, dtype=positions.dtype
            ),
        )
        selected_for_mask.flat_values.set_shape([None])

        # Construct the mask which is a boolean RT
        # Scatter 0's to positions that have been selector for deletion.
        update_values = tf.zeros_like(selected_for_mask.flat_values, "int32")
        update_indices = selected_for_mask.flat_values
        update_indices = tf.expand_dims(update_indices, -1)
        update_indices = tf.cast(update_indices, "int32")
        mask_flat = tf.ones_like(inputs.flat_values, dtype="int32")
        mask_flat = tf.tensor_scatter_nd_update(
            mask_flat, update_indices, update_values
        )
        mask = tf.cast(inputs.with_flat_values(mask_flat), "bool")

        inputs = tf.ragged.boolean_mask(inputs, mask)

        if unbatched:
            inputs = tf.squeeze(inputs, axis=0)

        return inputs
//  [python]
def _process_experience_global(self, experience):
    """Processes the experience in case the agent accepts only global features.

    In the experience coming from the replay buffer, the reward (and all other
    elements) have two batch dimensions `batch_size` and `time_steps`, where
    `time_steps` is the number of driver steps executed in each training loop.
    We flatten the tensors in order to reflect the effective batch size. Then,
    all the necessary processing on the observation is done, including splitting
    the action mask if it is present.

    Args:
      experience: An instance of trajectory. Every element in the trajectory has
        two batch dimensions.

    Returns:
      A tuple of reward, action, observation, and batch_size. All the outputs
        (except `batch_size`) have a single batch dimension of value
        `batch_size`.
    """
    reward, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.reward, self._time_step_spec.reward
    )
    observation, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.observation, self.training_data_spec.observation
    )
    action, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.action, self._action_spec
    )
    batch_size = tf.cast(
        tf.compat.dimension_value(tf.shape(reward)[0]), dtype=tf.int64
    )

    if self._observation_and_action_constraint_splitter is not None:
      observation, _ = self._observation_and_action_constraint_splitter(
          observation
      )
    if self._add_bias:
      # The bias is added via a constant 1 feature.
      observation = tf.concat(
          [observation, tf.ones([batch_size, 1], dtype=observation.dtype)],
          axis=1,
      )

    observation = tf.reshape(
        t.cast(observation, self._dtype), [batch_size, -1]
    )
    reward = t.cast(reward, self._dtype)

    return reward, action, observation, batch_size