jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def nap_project_multitap_filter(filt: tf.Tensor, alpha: float = 0.7,
                                iterations: int = 1,
                                allow_scaling: bool = False) -> tf.Tensor:
  """Near-all-pass projection for a multitap filter."""
  if filt.dtype.is_floating:
    forward_fft = tf.signal.rfft
    inverse_fft = tf.signal.irfft
  elif filt.dtype.is_complex:
    forward_fft = tf.signal.fft
    inverse_fft = tf.signal.ifft
  else:
    raise ValueError(f'Unsupported filter type {filt.dtype}')
  filter_len = _smart_dim(filt, -1)
  filter_fft_len = _enclosing_power_of_2(filter_len)
  for _ in range(iterations):
    filt_fft = forward_fft(_pad_to_length(filt, filter_fft_len))
    filt_fft_mag = tf.math.abs(filt_fft)
    if allow_scaling:
      mean_mag = tf.reduce_mean(filt_fft_mag)
    else:
      mean_mag = 1.0
    norm_alpha_min = mean_mag * alpha
    norm_alpha_max = mean_mag * 1.0/alpha
    filt_fft_mag_clipped = tf.clip_by_value(filt_fft_mag, norm_alpha_min,
                                            norm_alpha_max)
    filt_projected = (t.complex64(filt_fft_mag_clipped) *
                      tf.math.exp(tf.dtypes.complex(0.0, 1.0) *
                                  t.complex64(tf.math.angle(filt_fft))))
    filt = inverse_fft(filt_projected)[..., :filter_len]
  return filt
//  [python]
def nap_project_multitap_filter(filt: tf.Tensor, alpha: float = 0.7,
                                iterations: int = 1,
                                allow_scaling: bool = False) -> tf.Tensor:
  """Near-all-pass projection for a multitap filter."""
  if filt.dtype.is_floating:
    forward_fft = tf.signal.rfft
    inverse_fft = tf.signal.irfft
  elif filt.dtype.is_complex:
    forward_fft = tf.signal.fft
    inverse_fft = tf.signal.ifft
  else:
    raise ValueError(f'Unsupported filter type {filt.dtype}')
  filter_len = _smart_dim(filt, -1)
  filter_fft_len = _enclosing_power_of_2(filter_len)
  for _ in range(iterations):
    filt_fft = forward_fft(_pad_to_length(filt, filter_fft_len))
    filt_fft_mag = tf.math.abs(filt_fft)
    if allow_scaling:
      mean_mag = tf.reduce_mean(filt_fft_mag)
    else:
      mean_mag = 1.0
    norm_alpha_min = mean_mag * alpha
    norm_alpha_max = mean_mag * 1.0/alpha
    filt_fft_mag_clipped = tf.clip_by_value(filt_fft_mag, norm_alpha_min,
                                            norm_alpha_max)
    filt_projected = (tf.cast(filt_fft_mag_clipped, tf.complex64) *
                      tf.math.exp(tf.dtypes.complex(0.0, 1.0) *
                                  tf.cast(tf.math.angle(filt_fft),
                                          tf.complex64)))
    filt = inverse_fft(filt_projected)[..., :filter_len]
  return filt
//  [python]
def inverse_double_stft(signal_dstft: tf.Tensor, frame_length: int,
                        frame_step: int, meta_frame_length: int,
                        meta_frame_step: int, nfft: Optional[int] = None,
                        pad_begin: bool = True,
                        forward_window_fn: Callable[..., Any] = _sqrt_hann,
                        output_len: Optional[tf.Tensor] = None) -> tf.Tensor:
  """Inverse double STFT."""
  signal_ftt = tf.signal.ifft(signal_dstft)
  i_window = tf.signal.inverse_stft_window_fn(
      meta_frame_step, forward_window_fn=forward_window_fn)(
          meta_frame_length, dtype=tf.float32)
  i_window = t.cast(i_window, tf.complex64)
  signal_ftt = i_window * signal_ftt[..., :meta_frame_length]
  signal_ft = tf.signal.overlap_and_add(signal_ftt, meta_frame_step)
  if pad_begin:
    signal_ft = _clip_beginning(signal_ft, meta_frame_length - meta_frame_step)
  signal_tf = tf.einsum('...ft->...tf', signal_ft)
  return _istft(signal_tf, frame_length, frame_step, nfft, forward_window_fn,
                output_len=output_len)
//  [python]
def inverse_double_stft(signal_dstft: tf.Tensor, frame_length: int,
                        frame_step: int, meta_frame_length: int,
                        meta_frame_step: int, nfft: Optional[int] = None,
                        pad_begin: bool = True,
                        forward_window_fn: Callable[..., Any] = _sqrt_hann,
                        output_len: Optional[tf.Tensor] = None) -> tf.Tensor:
  """Inverse double STFT."""
  signal_ftt = tf.signal.ifft(signal_dstft)
  i_window = tf.signal.inverse_stft_window_fn(
      meta_frame_step, forward_window_fn=forward_window_fn)(
          meta_frame_length, dtype=tf.float32)
  i_window = tf.cast(i_window, tf.complex64)
  signal_ftt = i_window * signal_ftt[..., :meta_frame_length]
  signal_ft = tf.signal.overlap_and_add(signal_ftt, meta_frame_step)
  if pad_begin:
    signal_ft = _clip_beginning(signal_ft, meta_frame_length - meta_frame_step)
  signal_tf = tf.einsum('...ft->...tf', signal_ft)
  return _istft(signal_tf, frame_length, frame_step, nfft, forward_window_fn,
                output_len=output_len)
//  [python]
def tf_mask_tokens(
        self, inputs: Any, vocab_size, mask_token_id, special_tokens_mask: Optional[Any] = None
    ) -> Tuple[Any, Any]:
        """
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
        """
        import tensorflow as tf

        mask_token_id = tf.cast(mask_token_id, inputs.dtype)

        input_shape = tf.shape(inputs)
        # 1 for a special token, 0 for a normal token in the special tokens mask
        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
        masked_indices = self.tf_bernoulli(input_shape, self.mlm_probability) & ~special_tokens_mask
        # Replace unmasked indices with -100 in the labels since we only compute loss on masked tokens
        labels = tf.where(masked_indices, inputs, -100)

        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
        indices_replaced = self.tf_bernoulli(input_shape, 0.8) & masked_indices

        inputs = tf.where(indices_replaced, mask_token_id, inputs)

        # 10% of the time, we replace masked input tokens with random word
        indices_random = self.tf_bernoulli(input_shape, 0.1) & masked_indices & ~indices_replaced
        random_words = tf.random.uniform(input_shape, maxval=vocab_size, dtype=inputs.dtype)

        inputs = tf.where(indices_random, random_words, inputs)

        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
        return inputs, labels
//  [python]
def tf_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
        import tensorflow as tf

        # Handle dict or lists with proper padding and conversion to tensor.
        if isinstance(examples[0], Mapping):
            batch = pad_without_fast_tokenizer_warning(
                self.tokenizer, examples, return_tensors="tf", pad_to_multiple_of=self.pad_to_multiple_of
            )
        else:
            batch = {
                "input_ids": _tf_collate_batch(examples, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)
            }

        # If special token mask has been preprocessed, pop it from the dict.
        special_tokens_mask = batch.pop("special_tokens_mask", None)
        if self.mlm:
            if special_tokens_mask is None:
                special_tokens_mask = [
                    self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True)
                    for val in batch["input_ids"].numpy().tolist()
                ]
                # Cannot directly create as bool
                special_tokens_mask = tf.cast(tf.convert_to_tensor(special_tokens_mask, dtype=tf.int64), tf.bool)
            else:
                special_tokens_mask = tf.cast(special_tokens_mask, tf.bool)
            batch["input_ids"], batch["labels"] = self.tf_mask_tokens(
                tf.cast(batch["input_ids"], tf.int64),
                special_tokens_mask=special_tokens_mask,
                mask_token_id=self.tokenizer.mask_token_id,
                vocab_size=len(self.tokenizer),
            )
        else:
            labels = batch["input_ids"]
            if self.tokenizer.pad_token_id is not None:
                # Replace self.tokenizer.pad_token_id with -100
                labels = tf.where(labels == self.tokenizer.pad_token_id, -100, labels)
            else:
                labels = tf.identity(labels)  # Makes a copy, just in case
            batch["labels"] = labels
        return batch
//  [python]
def tf_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
        import tensorflow as tf

        if isinstance(examples[0], Mapping):
            input_ids = [e["input_ids"] for e in examples]
        else:
            input_ids = examples
            examples = [{"input_ids": e} for e in examples]

        batch_input = _tf_collate_batch(input_ids, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)

        mask_labels = []
        for e in examples:
            ref_tokens = []
            for id in tolist(e["input_ids"]):
                token = self.tokenizer._convert_id_to_token(id)
                ref_tokens.append(token)

            # For Chinese tokens, we need extra inf to mark sub-word, e.g [喜,欢]-> [喜，##欢]
            if "chinese_ref" in e:
                ref_pos = tolist(e["chinese_ref"])
                len_seq = len(e["input_ids"])
                for i in range(len_seq):
                    if i in ref_pos:
                        ref_tokens[i] = "##" + ref_tokens[i]
            mask_labels.append(self._whole_word_mask(ref_tokens))
        batch_mask = _tf_collate_batch(mask_labels, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)
        inputs, labels = self.tf_mask_tokens(tf.to_int64(batch_input), batch_mask)
        return {"input_ids": inputs, "labels": labels}
//  [python]
def tf_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
        import tensorflow as tf

        if isinstance(examples[0], Mapping):
            input_ids = [e["input_ids"] for e in examples]
        else:
            input_ids = examples
            examples = [{"input_ids": e} for e in examples]

        batch_input = _tf_collate_batch(input_ids, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)

        mask_labels = []
        for e in examples:
            ref_tokens = []
            for id in tolist(e["input_ids"]):
                token = self.tokenizer._convert_id_to_token(id)
                ref_tokens.append(token)

            # For Chinese tokens, we need extra inf to mark sub-word, e.g [喜,欢]-> [喜，##欢]
            if "chinese_ref" in e:
                ref_pos = tolist(e["chinese_ref"])
                len_seq = len(e["input_ids"])
                for i in range(len_seq):
                    if i in ref_pos:
                        ref_tokens[i] = "##" + ref_tokens[i]
            mask_labels.append(self._whole_word_mask(ref_tokens))
        batch_mask = _tf_collate_batch(mask_labels, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)
        inputs, labels = self.tf_mask_tokens(tf.cast(batch_input, tf.int64), batch_mask)
        return {"input_ids": inputs, "labels": labels}
//  [python]
def _build_dynamics_net(self, s, a, s_):
        with tf.variable_scope("dyn_net"):
            float_a = tf.expand_dims(tf.to_float(a, name="float_a"), axis=1, name="2d_a")
            sa = tf.concat((s, float_a), axis=1, name="sa")
            encoded_s_ = s_                # here we use s_ as the encoded s_

            dyn_l = tf.layers.dense(sa, 32, activation=tf.nn.relu)
            dyn_s_ = tf.layers.dense(dyn_l, self.n_s)  # predicted s_
        with tf.name_scope("int_r"):
            squared_diff = tf.reduce_sum(tf.square(encoded_s_ - dyn_s_), axis=1)  # intrinsic reward

        # It is better to reduce the learning rate in order to stay curious
        train_op = tf.train.RMSPropOptimizer(self.lr, name="dyn_opt").minimize(tf.reduce_mean(squared_diff))
        return dyn_s_, squared_diff, train_op
//  [python]
def _build_dynamics_net(self, s, a, s_):
        with tf.variable_scope("dyn_net"):
            float_a = tf.expand_dims(tf.cast(a, dtype=tf.float32, name="float_a"), axis=1, name="2d_a")
            sa = tf.concat((s, float_a), axis=1, name="sa")
            encoded_s_ = s_                # here we use s_ as the encoded s_

            dyn_l = tf.layers.dense(sa, 32, activation=tf.nn.relu)
            dyn_s_ = tf.layers.dense(dyn_l, self.n_s)  # predicted s_
        with tf.name_scope("int_r"):
            squared_diff = tf.reduce_sum(tf.square(encoded_s_ - dyn_s_), axis=1)  # intrinsic reward

        # It is better to reduce the learning rate in order to stay curious
        train_op = tf.train.RMSPropOptimizer(self.lr, name="dyn_opt").minimize(tf.reduce_mean(squared_diff))
        return dyn_s_, squared_diff, train_op
//  [python]
def tf_mask_tokens(self, inputs: Any, mask_labels: Any) -> Tuple[Any, Any]:
        """
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. Set
        'mask_labels' means we use whole word mask (wwm), we directly mask idxs according to it's ref.
        """
        import tensorflow as tf

        input_shape = tf.shape(inputs)
        if self.tokenizer.mask_token is None:
            raise ValueError(
                "This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the"
                " --mlm flag if you want to use this tokenizer."
            )
        labels = tf.identity(inputs)
        # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)

        masked_indices = tf.cast(mask_labels, tf.bool)

        special_tokens_mask = [
            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels
        ]
        masked_indices = masked_indices & ~tf.cast(special_tokens_mask, dtype=tf.bool)
        if self.tokenizer._pad_token is not None:
            padding_mask = inputs == self.tokenizer.pad_token_id
            masked_indices = masked_indices & ~padding_mask

        # Replace unmasked indices with -100 in the labels since we only compute loss on masked tokens
        labels = tf.where(masked_indices, inputs, -100)

        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
        indices_replaced = self.tf_bernoulli(input_shape, 0.8) & masked_indices

        inputs = tf.where(indices_replaced, self.tokenizer.mask_token_id, inputs)

        # 10% of the time, we replace masked input tokens with random word
        indices_random = self.tf_bernoulli(input_shape, 0.5) & masked_indices & ~indices_replaced
        random_words = tf.random.uniform(input_shape, maxval=len(self.tokenizer), dtype=tf.int64)
        inputs = tf.where(indices_random, random_words, inputs)

        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
        return inputs, labels
//  [python]
def _make_grid(nx=20, ny=20):
        """Generates a grid of shape [1, 1, ny * nx, 2] with ranges [0, nx) and [0, ny) for object detection."""
        # return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
        xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
        return tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]).astype(tf.float32)
//  [python]
def _make_grid(nx=20, ny=20):
        """Generates a grid of shape [1, 1, ny * nx, 2] with ranges [0, nx) and [0, ny) for object detection."""
        # return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
        xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
        return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32)
//  [python]
def project(self, img_feat, x, y, dim):
        x1 = tf.floor(x)
        x2 = tf.ceil(x)
        y1 = tf.floor(y)
        y2 = tf.ceil(y)
        Q11 = tf.gather_nd(img_feat, tf.stack([tf.to_int32(x1), tf.to_int32(y1)],1))
        Q12 = tf.gather_nd(img_feat, tf.stack([tf.to_int32(x1), tf.to_int32(y2)],1))
        Q21 = tf.gather_nd(img_feat, tf.stack([tf.to_int32(x2), tf.to_int32(y1)],1))
        Q22 = tf.gather_nd(img_feat, tf.stack([tf.to_int32(x2), tf.to_int32(y2)],1))

        weights = tf.multiply(tf.subtract(x2,x), tf.subtract(y2,y))
        Q11 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q11)

        weights = tf.multiply(tf.subtract(x,x1), tf.subtract(y2,y))
        Q21 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q21)

        weights = tf.multiply(tf.subtract(x2,x), tf.subtract(y,y1))
        Q12 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q12)

        weights = tf.multiply(tf.subtract(x,x1), tf.subtract(y,y1))
        Q22 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q22)

        outputs = tf.add_n([Q11, Q21, Q12, Q22])
        return outputs
//  [python]
def project(self, img_feat, x, y, dim):
        x1 = tf.floor(x)
        x2 = tf.ceil(x)
        y1 = tf.floor(y)
        y2 = tf.ceil(y)
        Q11 = tf.gather_nd(img_feat, tf.stack([tf.cast(x1,tf.int32), tf.cast(y1,tf.int32)],1))
        Q12 = tf.gather_nd(img_feat, tf.stack([tf.cast(x1,tf.int32), tf.cast(y2,tf.int32)],1))
        Q21 = tf.gather_nd(img_feat, tf.stack([tf.cast(x2,tf.int32), tf.cast(y1,tf.int32)],1))
        Q22 = tf.gather_nd(img_feat, tf.stack([tf.cast(x2,tf.int32), tf.cast(y2,tf.int32)],1))

        weights = tf.multiply(tf.subtract(x2,x), tf.subtract(y2,y))
        Q11 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q11)

        weights = tf.multiply(tf.subtract(x,x1), tf.subtract(y2,y))
        Q21 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q21)

        weights = tf.multiply(tf.subtract(x2,x), tf.subtract(y,y1))
        Q12 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q12)

        weights = tf.multiply(tf.subtract(x,x1), tf.subtract(y,y1))
        Q22 = tf.multiply(tf.tile(tf.reshape(weights,[-1,1]),[1,dim]), Q22)

        outputs = tf.add_n([Q11, Q21, Q12, Q22])
        return outputs
//  [python]
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):  # agnostic NMS
        """Performs non-max suppression on bounding boxes with class, IoU, and confidence thresholds; returns processed
        boxes, scores, classes, and count.
        """
        boxes, classes, scores = x
        class_inds = tf.to_float(tf.argmax(classes, axis=-1))
        scores_inp = tf.reduce_max(scores, -1)
        selected_inds = tf.image.non_max_suppression(
            boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres
        )
        selected_boxes = tf.gather(boxes, selected_inds)
        padded_boxes = tf.pad(
            selected_boxes,
            paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
            mode="CONSTANT",
            constant_values=0.0,
        )
        selected_scores = tf.gather(scores_inp, selected_inds)
        padded_scores = tf.pad(
            selected_scores,
            paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
            mode="CONSTANT",
            constant_values=-1.0,
        )
        selected_classes = tf.gather(class_inds, selected_inds)
        padded_classes = tf.pad(
            selected_classes,
            paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
            mode="CONSTANT",
            constant_values=-1.0,
        )
        valid_detections = tf.shape(selected_inds)[0]
        return padded_boxes, padded_scores, padded_classes, valid_detections
//  [python]
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):  # agnostic NMS
        """Performs non-max suppression on bounding boxes with class, IoU, and confidence thresholds; returns processed
        boxes, scores, classes, and count.
        """
        boxes, classes, scores = x
        class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
        scores_inp = tf.reduce_max(scores, -1)
        selected_inds = tf.image.non_max_suppression(
            boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres
        )
        selected_boxes = tf.gather(boxes, selected_inds)
        padded_boxes = tf.pad(
            selected_boxes,
            paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
            mode="CONSTANT",
            constant_values=0.0,
        )
        selected_scores = tf.gather(scores_inp, selected_inds)
        padded_scores = tf.pad(
            selected_scores,
            paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
            mode="CONSTANT",
            constant_values=-1.0,
        )
        selected_classes = tf.gather(class_inds, selected_inds)
        padded_classes = tf.pad(
            selected_classes,
            paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
            mode="CONSTANT",
            constant_values=-1.0,
        )
        valid_detections = tf.shape(selected_inds)[0]
        return padded_boxes, padded_scores, padded_classes, valid_detections
//  [python]
def mrcnn_class_loss_graph(target_class_ids, pred_class_logits, active_class_ids):
    """Loss for the classifier head of Mask RCNN.

    target_class_ids: [batch, num_rois]. Integer class IDs. Uses zero
        padding to fill in the array.
    pred_class_logits: [batch, num_rois, num_classes]
    active_class_ids: [batch, num_classes]. Has a value of 1 for
        classes that are in the dataset of the image, and 0
        for classes that are not in the dataset.
    """
    # During model building, Keras calls this function with
    # target_class_ids of type float32. Unclear why. Cast it
    # to int to get around it.
    target_class_ids = tf.to_int64(target_class_ids)

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(input=pred_class_logits, axis=2)
    #       images in a batch have the same active_class_ids
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    loss = tf.reduce_sum(input_tensor=loss) / tf.reduce_sum(input_tensor=pred_active)
    return loss
//  [python]
def mrcnn_class_loss_graph(target_class_ids, pred_class_logits, active_class_ids):
    """Loss for the classifier head of Mask RCNN.

    target_class_ids: [batch, num_rois]. Integer class IDs. Uses zero
        padding to fill in the array.
    pred_class_logits: [batch, num_rois, num_classes]
    active_class_ids: [batch, num_classes]. Has a value of 1 for
        classes that are in the dataset of the image, and 0
        for classes that are not in the dataset.
    """
    # During model building, Keras calls this function with
    # target_class_ids of type float32. Unclear why. Cast it
    # to int to get around it.
    target_class_ids = tf.cast(target_class_ids, 'int64')

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(input=pred_class_logits, axis=2)
    #       images in a batch have the same active_class_ids
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    loss = tf.reduce_sum(input_tensor=loss) / tf.reduce_sum(input_tensor=pred_active)
    return loss
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indices.
    positive_roi_ix = tf.compat.v1.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_roi_ix), tf.int64)
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(tf.size(input=target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def bilinear_resize_with_pad(
    images: tf.Tensor, rescale_size: tf.Tensor, output_size: tf.Tensor
) -> tf.Tensor:
  """Bilinear resizes the images, then pads to output size.

  Args:
    images: A tensor in shape (batch_size, input_h, input_w, ...) with arbitrary
      numbers of channel dimensions.
    rescale_size: An int tensor in shape (2,) or (batch_size, 2), representing
      the sizes of the rescaled images.
    output_size: The size of the output image in (output_h, output_w).

  Returns:
    A tensor in shape (batch_size, output_h, output_w, ...). The result has the
    same dtype as the input if it's float32, float16, bfloat16, otherwise the
    result is float32.
  """
  images_shape = images.get_shape().as_list()
  images_rank = len(images_shape)
  if images_rank < 3:
    raise ValueError(
        'Expected the input images (batch_size, height, width, ...) '
        'has rank >= 3, was: %s' % images_shape
    )
  batch_size = tf.shape(images)[0]
  rescale_size = tf.convert_to_tensor(rescale_size)
  if len(rescale_size.get_shape().as_list()) == 1:
    rescale_size = tf.broadcast_to(rescale_size, [batch_size, 2])

  # Rescales the images, applies the offset and pastes to the output canvas.

  # (batch_size, 2)
  ymin_xmin = tf.broadcast_to([0, 0], [batch_size, 2])
  # (batch_size, 2)
  ymax_xmax = tf.to_float(ymin_xmin) + rescale_size
  # (batch_size, 4)
  rescale_bbox = tf.concat([ymin_xmin, ymax_xmax], axis=1)
  # (batch_size, output_height, output_width, ...)
  return bilinear_resize_to_bbox(images, rescale_bbox, output_size)
//  [python]
def bilinear_resize_with_pad(
    images: tf.Tensor, rescale_size: tf.Tensor, output_size: tf.Tensor
) -> tf.Tensor:
  """Bilinear resizes the images, then pads to output size.

  Args:
    images: A tensor in shape (batch_size, input_h, input_w, ...) with arbitrary
      numbers of channel dimensions.
    rescale_size: An int tensor in shape (2,) or (batch_size, 2), representing
      the sizes of the rescaled images.
    output_size: The size of the output image in (output_h, output_w).

  Returns:
    A tensor in shape (batch_size, output_h, output_w, ...). The result has the
    same dtype as the input if it's float32, float16, bfloat16, otherwise the
    result is float32.
  """
  images_shape = images.get_shape().as_list()
  images_rank = len(images_shape)
  if images_rank < 3:
    raise ValueError(
        'Expected the input images (batch_size, height, width, ...) '
        'has rank >= 3, was: %s' % images_shape
    )
  batch_size = tf.shape(images)[0]
  rescale_size = tf.convert_to_tensor(rescale_size)
  if len(rescale_size.get_shape().as_list()) == 1:
    rescale_size = tf.broadcast_to(rescale_size, [batch_size, 2])

  # Rescales the images, applies the offset and pastes to the output canvas.

  # (batch_size, 2)
  ymin_xmin = tf.broadcast_to([0, 0], [batch_size, 2])
  # (batch_size, 2)
  ymax_xmax = tf.cast(ymin_xmin, rescale_size.dtype) + rescale_size
  # (batch_size, 4)
  rescale_bbox = tf.concat([ymin_xmin, ymax_xmax], axis=1)
  # (batch_size, output_height, output_width, ...)
  return bilinear_resize_to_bbox(images, rescale_bbox, output_size)
//  [python]
def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(input=target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(input=pred_masks)
    pred_masks = K.reshape(pred_masks, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(a=pred_masks, perm=[0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.compat.v1.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_ix))
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(input=y_true) > 0, K.binary_crossentropy(target=y_true, output=y_pred), tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(input=target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(input=pred_masks)
    pred_masks = K.reshape(pred_masks, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(a=pred_masks, perm=[0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.compat.v1.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(input=y_true) > 0, K.binary_crossentropy(target=y_true, output=y_pred), tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def call(self, inputs, training=None, mask=None):
    outputs = tf.nest.pack_sequence_as(inputs, tf.nest.flatten(inputs))
    for feature, type_str in self.config.features.items():
      assert type_str in self._type_map
      if type_str == "bfloat16":
        logging.warning(
          "Although bfloat16 and float32 have the same number of exponent bits, this down casting is not 100% lossless. Please double check metrics."
        )
      down_cast_data_type = self._type_map[type_str]
      outputs[feature] = t.cast(outputs[feature], dtype=down_cast_data_type)
    return outputs
//  [python]
def call(self, inputs, training=None, mask=None):
    outputs = tf.nest.pack_sequence_as(inputs, tf.nest.flatten(inputs))
    for feature, type_str in self.config.features.items():
      assert type_str in self._type_map
      if type_str == "bfloat16":
        logging.warning(
          "Although bfloat16 and float32 have the same number of exponent bits, this down casting is not 100% lossless. Please double check metrics."
        )
      down_cast_data_type = self._type_map[type_str]
      outputs[feature] = tf.cast(outputs[feature], dtype=down_cast_data_type)
    return outputs
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=self.config.GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(input_tensor=layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.to_float(tf.size(input=w))
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(input_tensor=layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=self.config.GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(input_tensor=layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.cast(tf.size(input=w), tf.float32)
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(input_tensor=layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def _create_model(self):
        # input layer to get embedding of features
        with tf.variable_scope('input_layer',
                               partitioner=self._input_layer_partitioner,
                               reuse=tf.AUTO_REUSE):
            uid_emb, item_emb, his_item_emb, sequence_length = self._embedding_input_layer(
            )

            item_his_eb_sum = tf.reduce_sum(his_item_emb, 1)
            mask = tf.sequence_mask(sequence_length)

        # Attention layer
        attention_scope = tf.variable_scope('attention_layer')
        with attention_scope.keep_weights(dtype=tf.float32) if self.bf16 \
            else attention_scope:
            attention_output = self._attention(item_emb, his_item_emb,
                                               self._attention_size, mask)
            att_fea = tf.reduce_sum(attention_output, 1)
            tf.summary.histogram('alpha_outputs', att_fea)

        top_input = tf.concat([
            uid_emb, item_emb, item_his_eb_sum, item_emb * item_his_eb_sum,
            att_fea
        ], 1)

        if self.bf16:
            top_input = tf.cast(top_input, tf.bfloat16)
        # Top MLP layer
        top_mlp_scope = tf.variable_scope(
            'top_mlp_layer',
            partitioner=self._dense_layer_partitioner,
            reuse=tf.AUTO_REUSE)
        with top_mlp_scope.keep_weights(dtype=tf.float32) if self.bf16 \
            else top_mlp_scope:
            self._logits = self._top_fc_layer(top_input)
        if self.bf16:
            self._logits = tf.cast(self._logits, dtype=tf.float32)

        self.probability = tf.math.sigmoid(self._logits)
        self.output = tf.round(self.probability)
//  [python]
def extract_subject(output, inputs):
        """根据subject_ids从output中取出subject的向量表征
        """
        subject_ids = inputs
        subject_ids = tf.to_int32(subject_ids)
        start = tf.batch_gather(output, subject_ids[:, :1])
        end = tf.batch_gather(output, subject_ids[:, 1:])
        subject = tf.concat([start, end], axis=2)
        return subject[:, 0]
//  [python]
def extract_subject(output, inputs):
        """根据subject_ids从output中取出subject的向量表征
        """
        subject_ids = inputs
        subject_ids = tf.cast(subject_ids, 'int32')
        start = tf.batch_gather(output, subject_ids[:, :1])
        end = tf.batch_gather(output, subject_ids[:, 1:])
        subject = tf.concat([start, end], axis=2)
        return subject[:, 0]
//  [python]
def compute_loss(
        model_output: tf.Tensor,
        gt: tf.Tensor,
        seq_len: List[int],
    ) -> tf.Tensor:
        """Compute categorical cross-entropy loss for the model.
        Sequences are masked after the EOS character.

        Args:
        ----
            model_output: predicted logits of the model
            gt: the encoded tensor with gt labels
            seq_len: lengths of each gt word inside the batch

        Returns:
        -------
            The loss of the model on the batch
        """
        # Input length : number of steps
        input_len = tf.shape(model_output)[1]
        # Add one for additional <eos> token (sos disappear in shift!)
        seq_len = tf.cast(seq_len, tf.int32) + 1
        # One-hot gt labels
        oh_gt = tf.one_hot(gt, depth=model_output.shape[2])
        # Compute loss: don't forget to shift gt! Otherwise the model learns to output the gt[t-1]!
        # The "masked" first gt char is <sos>.
        cce = tf.nn.softmax_cross_entropy_with_logits(oh_gt[:, 1:, :], model_output)
        # Compute mask
        mask_values = tf.zeros_like(cce)
        mask_2d = tf.sequence_mask(seq_len, input_len)
        masked_loss = tf.where(mask_2d, cce, mask_values)
        ce_loss = tf.math.divide(tf.reduce_sum(masked_loss, axis=1), tf.cast(seq_len, model_output.dtype))

        return tf.expand_dims(ce_loss, axis=1)
//  [python]
def call(self, inputs, training=None, mask=None):
    labels = self.config.engagements_list
    # union of engagements
    mask = tf.squeeze(tf.reduce_any(tf.stack([inputs[label] == 1 for label in labels], 1), 1))
    n_positives = tf.reduce_sum(t.cast(mask, t.int32))
    batch_size = t.cast(tf.shape(inputs[labels[0]])[0] / self.config.batch_multiplier, t.int32)
    negative_weights = tf.math.divide_no_nan(
      t.cast(self.config.batch_multiplier * batch_size - n_positives, t.float32),
      t.cast(batch_size - n_positives, t.float32),
    )
    new_weights = t.cast(mask, t.float32) + (1 - t.cast(mask, t.float32)) * negative_weights

    def _split_by_label_concatenate_and_truncate(input_tensor):
      # takes positive examples and concatenate with negative examples and truncate
      # DANGER: if n_positives > batch_size down-sampling is incorrect (do not use pb_50)
      return tf.concat(
        [
          input_tensor[mask],
          input_tensor[tf.math.logical_not(mask)],
        ],
        0,
      )[:batch_size]

    if "weights" not in inputs:
      # add placeholder so logic below applies even if weights aren't present in inputs
      inputs["weights"] = tf.ones([tf.shape(inputs[labels[0]])[0], self.config.num_engagements])

    for tensor in inputs:
      if tensor == "weights":
        inputs[tensor] = inputs[tensor] * tf.reshape(new_weights, [-1, 1])

      inputs[tensor] = _split_by_label_concatenate_and_truncate(inputs[tensor])

    return inputs
//  [python]
def call(self, inputs, training=None, mask=None):
    labels = self.config.engagements_list
    # union of engagements
    mask = tf.squeeze(tf.reduce_any(tf.stack([inputs[label] == 1 for label in labels], 1), 1))
    n_positives = tf.reduce_sum(tf.cast(mask, tf.int32))
    batch_size = tf.cast(tf.shape(inputs[labels[0]])[0] / self.config.batch_multiplier, tf.int32)
    negative_weights = tf.math.divide_no_nan(
      tf.cast(self.config.batch_multiplier * batch_size - n_positives, tf.float32),
      tf.cast(batch_size - n_positives, tf.float32),
    )
    new_weights = tf.cast(mask, tf.float32) + (1 - tf.cast(mask, tf.float32)) * negative_weights

    def _split_by_label_concatenate_and_truncate(input_tensor):
      # takes positive examples and concatenate with negative examples and truncate
      # DANGER: if n_positives > batch_size down-sampling is incorrect (do not use pb_50)
      return tf.concat(
        [
          input_tensor[mask],
          input_tensor[tf.math.logical_not(mask)],
        ],
        0,
      )[:batch_size]

    if "weights" not in inputs:
      # add placeholder so logic below applies even if weights aren't present in inputs
      inputs["weights"] = tf.ones([tf.shape(inputs[labels[0]])[0], self.config.num_engagements])

    for tensor in inputs:
      if tensor == "weights":
        inputs[tensor] = inputs[tensor] * tf.reshape(new_weights, [-1, 1])

      inputs[tensor] = _split_by_label_concatenate_and_truncate(inputs[tensor])

    return inputs
//  [python]
def _eval_metrics_fn():
    return {
        "acc": Accuracy(),
        "mse": MeanSquaredError(),
        "acc_fn": lambda labels, outputs: t.equal(
            t.cast(outputs, t.int32), t.cast(labels, t.int32)
        ),
    }
//  [python]
def _eval_metrics_fn():
    return {
        "acc": Accuracy(),
        "mse": MeanSquaredError(),
        "acc_fn": lambda labels, outputs: tf.equal(
            tf.cast(outputs, tf.int32), tf.cast(labels, tf.int32)
        ),
    }
//  [python]
def call(
        self,
        x: tf.Tensor,
        target: Optional[List[str]] = None,
        return_model_output: bool = False,
        return_preds: bool = False,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        features = self.feat_extractor(x, **kwargs)  # (batch_size, patches_seqlen, d_model)

        if target is not None:
            gt, seq_len = self.build_target(target)
            seq_len = t.cast(seq_len, tf.int32)

        if kwargs.get("training", False) and target is None:
            raise ValueError("Need to provide labels during training")

        features = features[:, : self.max_length]  # (batch_size, max_length, d_model)
        B, N, E = features.shape
        features = tf.reshape(features, (B * N, E))
        logits = tf.reshape(
            self.head(features, **kwargs), (B, N, len(self.vocab) + 1)
        )  # (batch_size, max_length, vocab + 1)
        decoded_features = _bf16_to_float32(logits[:, 1:])  # remove cls_token

        out: Dict[str, tf.Tensor] = {}
        if self.exportable:
            out["logits"] = decoded_features
            return out

        if return_model_output:
            out["out_map"] = decoded_features

        if target is None or return_preds:
            # Post-process boxes
            out["preds"] = self.postprocessor(decoded_features)

        if target is not None:
            out["loss"] = self.compute_loss(decoded_features, gt, seq_len)

        return out
//  [python]
def call(
        self,
        x: tf.Tensor,
        target: Optional[List[str]] = None,
        return_model_output: bool = False,
        return_preds: bool = False,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        features = self.feat_extractor(x, **kwargs)  # (batch_size, patches_seqlen, d_model)

        if target is not None:
            gt, seq_len = self.build_target(target)
            seq_len = tf.cast(seq_len, tf.int32)

        if kwargs.get("training", False) and target is None:
            raise ValueError("Need to provide labels during training")

        features = features[:, : self.max_length]  # (batch_size, max_length, d_model)
        B, N, E = features.shape
        features = tf.reshape(features, (B * N, E))
        logits = tf.reshape(
            self.head(features, **kwargs), (B, N, len(self.vocab) + 1)
        )  # (batch_size, max_length, vocab + 1)
        decoded_features = _bf16_to_float32(logits[:, 1:])  # remove cls_token

        out: Dict[str, tf.Tensor] = {}
        if self.exportable:
            out["logits"] = decoded_features
            return out

        if return_model_output:
            out["out_map"] = decoded_features

        if target is None or return_preds:
            # Post-process boxes
            out["preds"] = self.postprocessor(decoded_features)

        if target is not None:
            out["loss"] = self.compute_loss(decoded_features, gt, seq_len)

        return out
//  [python]
def focal_loss_softmax(labels, logits, gamma=2):
    """
    https://github.com/fudannlp16/focal-loss/blob/master/focal_loss.py
    Computer focal loss for multi classification
    Args:
      labels: A int32 tensor of shape [batch_size].
      logits: A float32 tensor of shape [batch_size,num_classes].
      gamma: A scalar for focal loss gamma hyper-parameter.
    Returns:
      A tensor of the same shape as `lables`
    """
    y_pred=tf.nn.softmax(logits,axis=-1) # [batch_size,num_classes]
    cross_ent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=tf.squeeze(labels, axis=1), logits=logits)
    cross_ent = tf.expand_dims(cross_ent, 1)
    labels = tf.to_int32(labels)
    L=((1.0-tf.batch_gather(y_pred, labels))**gamma)*cross_ent
    return L