jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, inputs, y_true, y_pred, training=None):
        """
        Define a masked rank 1 ListNet loss.
        This loss is useful for multi-label classification when we have multiple
        click labels per document. This is because the loss breaks down the comparison
        between y_pred and y_true into individual binary assessments.
        Ref -> https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/tr-2007-40.pdf

        Parameters
        ----------
        inputs: dict of dict of tensors
            Dictionary of input feature tensors
        y_true: tensor
            True labels
        y_pred: tensor
            Predicted scores
        training: boolean
            Boolean indicating whether the layer is being used in training mode

        Returns
        -------
        tensor
            Scalar sigmoid cross entropy loss tensor

        Notes
        -----
            Uses `mask` field to exclude padded records from contributing
            to the loss
        """
        mask = tf.cast(inputs[FeatureTypeKey.MASK], y_pred.dtype)
        y_true = tf.cast(y_true, y_pred.dtype)
        batch_size = tf.cast(tf.shape(y_true)[0], tf.float32)

        # Mask the padded records
        y_true = tf.gather_nd(y_true, tf.where(tf.equal(mask, tf.constant(1.0))))
        y_pred = tf.gather_nd(y_pred, tf.where(tf.equal(mask, tf.constant(1.0))))

        # Reshape the tensors so that we sum the losses from each record
        y_true = tf.expand_dims(tf.squeeze(y_true), axis=-1)
        y_pred = tf.expand_dims(tf.squeeze(y_pred), axis=-1)

        # Scale the sum of losses down by number of queries in the batch
        return tf.math.divide(self.loss_fn(y_true, y_pred), batch_size)
//  [python]
def yolo_head(
            self,
            feats,
            anchors,
            num_classes,
            input_shape,
            training=True
    ):

        num_anchors = len(anchors)
        anchors_tensor = tf.reshape(
            tf.constant(
                anchors, dtype=tf.float32), [
                1, 1, 1, num_anchors, 2])
        grid_size = tf.shape(feats)[1:3]
        predictions = tf.reshape(
            feats,
            [-1, grid_size[0], grid_size[1], num_anchors, num_classes + 5])
        grid_y = tf.tile(tf.reshape(
            tf.range(grid_size[0]), [-1, 1, 1, 1]), [1, grid_size[1], 1, 1])
        grid_x = tf.tile(tf.reshape(tf.range(grid_size[1]), [
            1, -1, 1, 1]), [grid_size[0], 1, 1, 1])
        grid = tf.concat([grid_x, grid_y], axis=-1)
        grid = tf.to_float(grid)
        box_xy = (tf.sigmoid(predictions[..., :2]) +
                  grid) / tf.to_float(grid_size[::-1])
        box_wh = tf.exp(
            predictions[..., 2:4]) * anchors_tensor / tf.to_float(input_shape[::-1])
        box_confidence = tf.sigmoid(predictions[..., 4:5])
        box_class_probs = tf.sigmoid(predictions[..., 5:])
        if training:
            return grid, predictions, box_xy, box_wh
        return box_xy, box_wh, box_confidence, box_class_probs
//  [python]
def yolo_head(
            self,
            feats,
            anchors,
            num_classes,
            input_shape,
            training=True
    ):

        num_anchors = len(anchors)
        anchors_tensor = tf.reshape(
            tf.constant(
                anchors, dtype=tf.float32), [
                1, 1, 1, num_anchors, 2])
        grid_size = tf.shape(feats)[1:3]
        predictions = tf.reshape(
            feats,
            [-1, grid_size[0], grid_size[1], num_anchors, num_classes + 5])
        grid_y = tf.tile(tf.reshape(
            tf.range(grid_size[0]), [-1, 1, 1, 1]), [1, grid_size[1], 1, 1])
        grid_x = tf.tile(tf.reshape(tf.range(grid_size[1]), [
            1, -1, 1, 1]), [grid_size[0], 1, 1, 1])
        grid = tf.concat([grid_x, grid_y], axis=-1)
        grid = tf.cast(grid, tf.float32)
        box_xy = (tf.sigmoid(predictions[..., :2]) +
                  grid) / tf.cast(grid_size[::-1], tf.float32)
        box_wh = tf.exp(
            predictions[..., 2:4]) * anchors_tensor / input_shape[::-1]
        box_confidence = tf.sigmoid(predictions[..., 4:5])
        box_class_probs = tf.sigmoid(predictions[..., 5:])
        if training:
            return grid, predictions, box_xy, box_wh
        return box_xy, box_wh, box_confidence, box_class_probs
//  [python]
def yolo_boxes_scores(
            self,
            feats,
            anchors,
            num_classes,
            input_shape,
            image_shape):
        input_shape = tf.to_float(input_shape)
        image_shape = tf.to_float(image_shape)
        box_xy, box_wh, box_confidence, box_class_probs = self.yolo_head(
            feats, anchors, num_classes, input_shape, training=False)
        box_yx = box_xy[..., ::-1]
        box_hw = box_wh[..., ::-1]
        new_shape = tf.round(
            image_shape *
            tf.reduce_min(
                input_shape /
                image_shape))
        offset = (input_shape - new_shape) / 2. / input_shape
        scale = input_shape / new_shape
        box_yx = (box_yx - offset) * scale
        box_hw = box_hw * scale

        box_min = box_yx - box_hw / 2.
        box_max = box_yx + box_hw / 2.
        boxes = tf.concat(
            [box_min[..., 0:1],
             box_min[..., 1:2],
             box_max[..., 0:1],
             box_max[..., 1:2]],
            axis=-1
        )
        boxes *= tf.concat([image_shape, image_shape], axis=-1)
        boxes = tf.reshape(boxes, [-1, 4])
        boxes_scores = box_confidence * box_class_probs
        boxes_scores = tf.reshape(boxes_scores, [-1, num_classes])
        return boxes, boxes_scores
//  [python]
def yolo_boxes_scores(
            self,
            feats,
            anchors,
            num_classes,
            input_shape,
            image_shape):
        input_shape = tf.cast(input_shape, tf.float32)
        image_shape = tf.cast(image_shape, tf.float32)
        box_xy, box_wh, box_confidence, box_class_probs = self.yolo_head(
            feats, anchors, num_classes, input_shape, training=False)
        box_yx = box_xy[..., ::-1]
        box_hw = box_wh[..., ::-1]
        new_shape = tf.round(
            image_shape *
            tf.reduce_min(
                input_shape /
                image_shape))
        offset = (input_shape - new_shape) / 2. / input_shape
        scale = input_shape / new_shape
        box_yx = (box_yx - offset) * scale
        box_hw = box_hw * scale

        box_min = box_yx - box_hw / 2.
        box_max = box_yx + box_hw / 2.
        boxes = tf.concat(
            [box_min[..., 0:1],
             box_min[..., 1:2],
             box_max[..., 0:1],
             box_max[..., 1:2]],
            axis=-1
        )
        boxes *= tf.concat([image_shape, image_shape], axis=-1)
        boxes = tf.reshape(boxes, [-1, 4])
        boxes_scores = box_confidence * box_class_probs
        boxes_scores = tf.reshape(boxes_scores, [-1, num_classes])
        return boxes, boxes_scores
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    logits = t.cast(model_outputs['logits'], tf.float32)
    masked_labels, masked_weights = _masked_labels_and_weights(labels)
    loss = tf_keras.losses.sparse_categorical_crossentropy(
        masked_labels, logits, from_logits=True)
    numerator_loss = tf.reduce_sum(loss * masked_weights)
    denominator_loss = tf.reduce_sum(masked_weights)
    loss = tf.math.divide_no_nan(numerator_loss, denominator_loss)
    return loss
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    logits = tf.cast(model_outputs['logits'], tf.float32)
    masked_labels, masked_weights = _masked_labels_and_weights(labels)
    loss = tf_keras.losses.sparse_categorical_crossentropy(
        masked_labels, logits, from_logits=True)
    numerator_loss = tf.reduce_sum(loss * masked_weights)
    denominator_loss = tf.reduce_sum(masked_weights)
    loss = tf.math.divide_no_nan(numerator_loss, denominator_loss)
    return loss
//  [python]
def _masked_labels_and_weights(y_true):
  """Masks negative values from token level labels.

  Args:
    y_true: Token labels, typically shape (batch_size, seq_len), where tokens
      with negative labels should be ignored during loss/accuracy calculation.

  Returns:
    (masked_y_true, masked_weights) where `masked_y_true` is the input
    with each negative label replaced with zero and `masked_weights` is 0.0
    where negative labels were replaced and 1.0 for original labels.
  """
  # Ignore the classes of tokens with negative values.
  mask = tf.greater_equal(y_true, 0)
  # Replace negative labels, which are out of bounds for some loss functions,
  # with zero.
  masked_y_true = tf.where(mask, y_true, 0)
  return masked_y_true, tf.to_float(mask)
//  [python]
def _masked_labels_and_weights(y_true):
  """Masks negative values from token level labels.

  Args:
    y_true: Token labels, typically shape (batch_size, seq_len), where tokens
      with negative labels should be ignored during loss/accuracy calculation.

  Returns:
    (masked_y_true, masked_weights) where `masked_y_true` is the input
    with each negative label replaced with zero and `masked_weights` is 0.0
    where negative labels were replaced and 1.0 for original labels.
  """
  # Ignore the classes of tokens with negative values.
  mask = tf.greater_equal(y_true, 0)
  # Replace negative labels, which are out of bounds for some loss functions,
  # with zero.
  masked_y_true = tf.where(mask, y_true, 0)
  return masked_y_true, tf.cast(mask, tf.float32)
//  [python]
def metric_fn(logits, dup_mask, match_mlperf):
  dup_mask = t.cast(dup_mask, t.float32)
  logits = tf.slice(logits, [0, 1], [-1, -1])
  in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
      logits, dup_mask, match_mlperf)
  metric_weights = t.cast(metric_weights, t.float32)
  return in_top_k, metric_weights
//  [python]
def metric_fn(logits, dup_mask, match_mlperf):
  dup_mask = tf.cast(dup_mask, tf.float32)
  logits = tf.slice(logits, [0, 1], [-1, -1])
  in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
      logits, dup_mask, match_mlperf)
  metric_weights = tf.cast(metric_weights, tf.float32)
  return in_top_k, metric_weights
//  [python]
def eval(aligned_images, model_path):
    with tf.Graph().as_default():
        sess = tf.Session()
        images_pl = tf.placeholder(tf.float32, shape=[None, 160, 160, 3], name='input_image')
        images = tf.map_fn(lambda frame: tf.reverse_v2(frame, [-1]), images_pl) #BGR TO RGB
        images_norm = tf.map_fn(lambda frame: tf.image.per_image_standardization(frame), images)
        train_mode = tf.placeholder(tf.bool)
        age_logits, gender_logits, _ = inception_resnet_v1.inference(images_norm, keep_probability=0.8,
                                                                     phase_train=train_mode,
                                                                     weight_decay=1e-5)
        gender = tf.argmax(tf.nn.softmax(gender_logits), 1)
        age_ = tf.constant([i for i in range(0, 101)])
        age = tf.reduce_sum(tf.multiply(tf.nn.softmax(age_logits), tf.cast(age_, tf.float32)), axis=1)
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())
        sess.run(init_op)
        saver = tf.train.Saver()
        ckpt = tf.train.get_checkpoint_state(model_path)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            print("restore and continue training!")
        else:
            pass
        return sess.run([age, gender], feed_dict={images_pl: aligned_images, train_mode: False})
//  [python]
def eval(aligned_images, model_path):
    with tf.Graph().as_default():
        sess = tf.Session()
        images_pl = tf.placeholder(tf.float32, shape=[None, 160, 160, 3], name='input_image')
        images = tf.map_fn(lambda frame: tf.reverse_v2(frame, [-1]), images_pl) #BGR TO RGB
        images_norm = tf.map_fn(lambda frame: tf.image.per_image_standardization(frame), images)
        train_mode = tf.placeholder(tf.bool)
        age_logits, gender_logits, _ = inception_resnet_v1.inference(images_norm, keep_probability=0.8,
                                                                     phase_train=train_mode,
                                                                     weight_decay=1e-5)
        gender = tf.argmax(tf.nn.softmax(gender_logits), 1)
        age_ = tf.cast(tf.constant([i for i in range(0, 101)]), tf.float32)
        age = tf.reduce_sum(tf.multiply(tf.nn.softmax(age_logits), age_), axis=1)
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())
        sess.run(init_op)
        saver = tf.train.Saver()
        ckpt = tf.train.get_checkpoint_state(model_path)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            print("restore and continue training!")
        else:
            pass
        return sess.run([age, gender], feed_dict={images_pl: aligned_images, train_mode: False})
//  [python]
def _prepare_model_inputs(self, inputs: Optional[tf.Tensor] = None, bos_token_id: Optional[int] = None):
        # TODO(Patrick) - adapt this function when making `generate` more flexible
        # for all kinds of input types
        if inputs is None:
            # if no `inputs` are passed create prompt of size (1,1) filled with BOS token
            if not isinstance(bos_token_id, int) or bos_token_id < 0:
                raise ValueError(
                    "you should either supply a context to complete as `input_ids` input "
                    "or a `bos_token_id` (integer >= 0) as a first token to start the generation."
                )
            return tf.fill((1, 1), bos_token_id).astype(tf.int32)

        return inputs
//  [python]
def _prepare_model_inputs(self, inputs: Optional[tf.Tensor] = None, bos_token_id: Optional[int] = None):
        # TODO(Patrick) - adapt this function when making `generate` more flexible
        # for all kinds of input types
        if inputs is None:
            # if no `inputs` are passed create prompt of size (1,1) filled with BOS token
            if not isinstance(bos_token_id, int) or bos_token_id < 0:
                raise ValueError(
                    "you should either supply a context to complete as `input_ids` input "
                    "or a `bos_token_id` (integer >= 0) as a first token to start the generation."
                )
            return tf.cast(tf.fill((1, 1), bos_token_id), dtype=tf.int32)

        return inputs
//  [python]
def _make_grid(nx=20, ny=20):
        """Generates a 2D grid of coordinates in (x, y) format with shape [1, 1, ny*nx, 2]."""
        # 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 2D grid of coordinates in (x, y) format with shape [1, 1, ny*nx, 2]."""
        # 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 build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    label_ids = labels[self.label_field]
    if self.task_config.model.num_classes == 1:
      loss = tf_keras.losses.mean_squared_error(label_ids, model_outputs)
    else:
      loss = tf_keras.losses.sparse_categorical_crossentropy(
          label_ids, tf.cast(model_outputs, tf.float32), from_logits=True)

    if aux_losses:
      loss += tf.add_n(aux_losses)
    return tf_utils.safe_mean(loss)
//  [python]
def WeakSupLoss(self):
        '''
        Define additional losses for weakly supervised segmentation
        Inputs:
            self.logits = tf.reshape(self.logits, [-1, config.num_classes])
            self.labels = tf.reshape(self.labels, [-1])
            self.mask_boolen = tf.reshape(self.mask_boolen, [-1])
        Returns:

        '''
        
        ## MIL Branch
        one_hot_labels = tf.one_hot(self.batch_labels, depth=self.config.num_classes, axis=-1)
        batch_size = int(tf.shape(self.batch_labels)[0]//2)
#         one_hot_labels = tf.reshape(one_hot_labels, [self.inputs.shape[0],self.inputs.shape[1],self.config.num_classes])
        L_gt = tf.to_float(tf.reduce_max(one_hot_labels[:batch_size], axis=1))  # inexact labels B*13
        L_pred = tf.reduce_max(self.batch_logits[:batch_size], axis=1)
        loss_ineaxct = tf.nn.sigmoid_cross_entropy_with_logits(labels=L_gt, logits=L_pred)  # B*K, independent loss for each class
        self.loss_inexact = 1* tf.reduce_mean(loss_ineaxct)
        
        ## Siamese Branch
        self.prob_logits = tf.nn.softmax(self.batch_logits, axis=-1)
        self.loss_siamese = 1* tf.reduce_mean(tf.reduce_sum((self.prob_logits[:batch_size] - self.prob_logits[batch_size:]) ** 2, axis=-1))
        
#         pdb.set_trace()
        ## Smooth Branch
#         self.loss_smooth = 10* SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain(self.prob_logits[0::2], self.inputs['features'][0::2, :, 0:3], M=self.batch_masks[0::2], Y=self.batch_labels[0::2], gamma=1e-1)
        self.loss_smooth = 10 * SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain_mustLink(self.prob_logits[:batch_size], self.inputs['features'][:batch_size, :, 0:6], M=self.batch_masks[:batch_size], Y=self.batch_labels[:batch_size], gamma=1e-1)
//  [python]
def WeakSupLoss(self):
        '''
        Define additional losses for weakly supervised segmentation
        Inputs:
            self.logits = tf.reshape(self.logits, [-1, config.num_classes])
            self.labels = tf.reshape(self.labels, [-1])
            self.mask_boolen = tf.reshape(self.mask_boolen, [-1])
        Returns:

        '''
        
        ## MIL Branch
        one_hot_labels = tf.one_hot(self.batch_labels, depth=self.config.num_classes, axis=-1)
        batch_size = int(tf.shape(self.batch_labels)[0]//2)
#         one_hot_labels = tf.reshape(one_hot_labels, [self.inputs.shape[0],self.inputs.shape[1],self.config.num_classes])
        L_gt = tf.cast(tf.reduce_max(one_hot_labels[:batch_size], axis=1), tf.float32)  # inexact labels B*13
        L_pred = tf.reduce_max(self.batch_logits[:batch_size], axis=1)
        loss_ineaxct = tf.nn.sigmoid_cross_entropy_with_logits(labels=L_gt, logits=L_pred)  # B*K, independent loss for each class
        self.loss_inexact = 1* tf.reduce_mean(loss_ineaxct)
        
        ## Siamese Branch
        self.prob_logits = tf.nn.softmax(self.batch_logits, axis=-1)
        self.loss_siamese = 1* tf.reduce_mean(tf.reduce_sum((self.prob_logits[:batch_size] - self.prob_logits[batch_size:]) ** 2, axis=-1))
        
#         pdb.set_trace()
        ## Smooth Branch
#         self.loss_smooth = 10* SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain(self.prob_logits[0::2], self.inputs['features'][0::2, :, 0:3], M=self.batch_masks[0::2], Y=self.batch_labels[0::2], gamma=1e-1)
        self.loss_smooth = 10 * SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain_mustLink(self.prob_logits[:batch_size], self.inputs['features'][:batch_size, :, 0:6], M=self.batch_masks[:batch_size], Y=self.batch_labels[:batch_size], gamma=1e-1)
//  [python]
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):
        """Performs agnostic non-maximum suppression (NMS) on detected objects, filtering based on IoU and confidence
        thresholds.
        """
        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):
        """Performs agnostic non-maximum suppression (NMS) on detected objects, filtering based on IoU and confidence
        thresholds.
        """
        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 train(self):
        y = tf.one_hot(self.y, args.class_size)
        loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=self.logits)
        self.loss = tf.reduce_sum(loss)
        self.train_op = tf.train.AdamOptimizer(args.learning_rate).minimize(self.loss)
        self.predict = tf.argmax(self.logits, axis=1, name="predictions")
        correct_prediction = tf.equal(tf.to_int32(self.predict), self.y)
        self.accuracy = tf.reduce_mean(tf.to_float(correct_prediction), name="Accuracy")
//  [python]
def train(self):
        y = tf.one_hot(self.y, args.class_size)
        loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=self.logits)
        self.loss = tf.reduce_sum(loss)
        self.train_op = tf.train.AdamOptimizer(args.learning_rate).minimize(self.loss)
        self.predict = tf.argmax(self.logits, axis=1, name="predictions")
        correct_prediction = tf.equal(tf.cast(self.predict, tf.int32), self.y)
        self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32), name="Accuracy")
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: Optional[tf.Tensor] = None,
        token_type_ids: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
        inputs_embeds: Optional[tf.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        inputs = input_values_processing(
            func=self.call,
            config=self.config,
            input_values=input_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )

        extract_features = self.feature_extractor(
            tf.to_float(inputs["input_values"]), training=inputs["training"]
        )
        # extract_features = tf.transpose(extract_features, perm=(0, 2, 1))

        if inputs["attention_mask"] is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(inputs["attention_mask"], -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(extract_features)[1], dtype=extract_features.dtype
            )

        hidden_states, extract_features = self.feature_projection(extract_features, training=inputs["training"])

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if inputs["training"]:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=inputs["output_attentions"],
            output_hidden_states=inputs["output_hidden_states"],
            return_dict=inputs["return_dict"],
            training=inputs["training"],
        )
        hidden_states = encoder_outputs[0]

        if not inputs["return_dict"]:
            return (hidden_states, extract_features) + encoder_outputs[1:]

        return TFWav2Vec2BaseModelOutput(
            last_hidden_state=hidden_states,
            extract_features=extract_features,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: Optional[tf.Tensor] = None,
        token_type_ids: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
        inputs_embeds: Optional[tf.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        inputs = input_values_processing(
            func=self.call,
            config=self.config,
            input_values=input_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )

        extract_features = self.feature_extractor(
            tf.cast(inputs["input_values"], tf.float32), training=inputs["training"]
        )
        # extract_features = tf.transpose(extract_features, perm=(0, 2, 1))

        if inputs["attention_mask"] is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(inputs["attention_mask"], -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(extract_features)[1], dtype=extract_features.dtype
            )

        hidden_states, extract_features = self.feature_projection(extract_features, training=inputs["training"])

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if inputs["training"]:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=inputs["output_attentions"],
            output_hidden_states=inputs["output_hidden_states"],
            return_dict=inputs["return_dict"],
            training=inputs["training"],
        )
        hidden_states = encoder_outputs[0]

        if not inputs["return_dict"]:
            return (hidden_states, extract_features) + encoder_outputs[1:]

        return TFWav2Vec2BaseModelOutput(
            last_hidden_state=hidden_states,
            extract_features=extract_features,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def _extract_features_batch(self, serialized_batch):
        features = tf.parse_example(
            serialized_batch,
            features={'images': tf.FixedLenFeature([], tf.string),
                'imagepaths': tf.FixedLenFeature([], tf.string),
                'labels': tf.VarLenFeature(tf.int64),
                 })

        bs = features['images'].shape[0]
        images = tf.decode_raw(features['images'], tf.uint8)
        w, h = tuple(CFG.ARCH.INPUT_SIZE)
        images = tf.cast(x=images, dtype=tf.float32)
        #images = tf.subtract(tf.divide(images, 128.0), 1.0)
        images = tf.reshape(images, [bs, h, -1, CFG.ARCH.INPUT_CHANNELS])

        labels = features['labels']
        labels = tf.to_int32(labels)

        imagepaths = features['imagepaths']

        return images, labels, imagepaths
//  [python]
def _extract_features_batch(self, serialized_batch):
        features = tf.parse_example(
            serialized_batch,
            features={'images': tf.FixedLenFeature([], tf.string),
                'imagepaths': tf.FixedLenFeature([], tf.string),
                'labels': tf.VarLenFeature(tf.int64),
                 })

        bs = features['images'].shape[0]
        images = tf.decode_raw(features['images'], tf.uint8)
        w, h = tuple(CFG.ARCH.INPUT_SIZE)
        images = tf.cast(x=images, dtype=tf.float32)
        #images = tf.subtract(tf.divide(images, 128.0), 1.0)
        images = tf.reshape(images, [bs, h, -1, CFG.ARCH.INPUT_CHANNELS])

        labels = features['labels']
        labels = tf.cast(labels, tf.int32)

        imagepaths = features['imagepaths']

        return images, labels, imagepaths
//  [python]
def process_source_id(source_id: tf.Tensor) -> tf.Tensor:
  """Processes source_id to the right format.

  Args:
    source_id: A `tf.Tensor` that contains the source ID. It can be empty.

  Returns:
    A formatted source ID.
  """
  if source_id.dtype == tf.string:
    source_id = tf.strings.to_number(source_id, tf.int64)
  with tf.control_dependencies([source_id]):
    source_id = tf.cond(
        pred=tf.equal(tf.size(input=source_id), 0),
        true_fn=lambda: tf.to_int64(tf.constant(-1)),
        false_fn=lambda: tf.identity(source_id))
  return source_id
//  [python]
def process_source_id(source_id: tf.Tensor) -> tf.Tensor:
  """Processes source_id to the right format.

  Args:
    source_id: A `tf.Tensor` that contains the source ID. It can be empty.

  Returns:
    A formatted source ID.
  """
  if source_id.dtype == tf.string:
    source_id = tf.strings.to_number(source_id, tf.int64)
  with tf.control_dependencies([source_id]):
    source_id = tf.cond(
        pred=tf.equal(tf.size(input=source_id), 0),
        true_fn=lambda: tf.cast(tf.constant(-1), tf.int64),
        false_fn=lambda: tf.identity(source_id))
  return source_id
//  [python]
def _shift_right(self, input_ids):
        decoder_start_token_id = self.config.decoder_start_token_id
        pad_token_id = self.config.pad_token_id

        assert decoder_start_token_id is not None, (
            "self.model.config.decoder_start_token_id has to be defined. In TF T5 it is usually set to the"
            " pad_token_id. See T5 docs for more information"
        )

        start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
        start_tokens = tf.to_int32(start_tokens)  # Ensure compatible dtypes for concatenation
        shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)

        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100,
            tf.to_int32(tf.fill(shape_list(shifted_input_ids), pad_token_id)),
            shifted_input_ids,
        )

        # "Verify that `labels` has only positive values and -100"
        assert_gte0 = tf.debugging.assert_greater_equal(
            shifted_input_ids, tf.constant(0, dtype=shifted_input_ids.dtype)
        )

        # Make sure the assertion op is called by wrapping the result in an identity no-op
        with tf.control_dependencies([assert_gte0]):
            shifted_input_ids = tf.identity(shifted_input_ids)

        return shifted_input_ids
//  [python]
def _shift_right(self, input_ids):
        decoder_start_token_id = self.config.decoder_start_token_id
        pad_token_id = self.config.pad_token_id

        assert decoder_start_token_id is not None, (
            "self.model.config.decoder_start_token_id has to be defined. In TF T5 it is usually set to the"
            " pad_token_id. See T5 docs for more information"
        )

        start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
        start_tokens = tf.cast(start_tokens, input_ids.dtype)  # Ensure compatible dtypes for concatenation
        shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)

        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100,
            tf.cast(tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids.dtype),
            shifted_input_ids,
        )

        # "Verify that `labels` has only positive values and -100"
        assert_gte0 = tf.debugging.assert_greater_equal(
            shifted_input_ids, tf.constant(0, dtype=shifted_input_ids.dtype)
        )

        # Make sure the assertion op is called by wrapping the result in an identity no-op
        with tf.control_dependencies([assert_gte0]):
            shifted_input_ids = tf.identity(shifted_input_ids)

        return shifted_input_ids
//  [python]
def binarize_matting_map(matting_map: tf.Tensor,
                         threshold: float = 0.5) -> tf.Tensor:
  """Binarizes a matting map.

  If the matting_map value is above a threshold, set it as 1 otherwise 0. The
  binarization is done for every element in the matting_map.

  Args:
    matting_map: The groundtruth in the matting map format.
    threshold: The threshold used to binarize the matting map.

  Returns:
    The binarized labels (0 for BG, 1 for FG) as tf.float32.
  """
  return tf.to_float(tf.greater(matting_map, threshold))
//  [python]
def binarize_matting_map(matting_map: tf.Tensor,
                         threshold: float = 0.5) -> tf.Tensor:
  """Binarizes a matting map.

  If the matting_map value is above a threshold, set it as 1 otherwise 0. The
  binarization is done for every element in the matting_map.

  Args:
    matting_map: The groundtruth in the matting map format.
    threshold: The threshold used to binarize the matting map.

  Returns:
    The binarized labels (0 for BG, 1 for FG) as tf.float32.
  """
  return tf.cast(tf.greater(matting_map, threshold), tf.float32)
//  [python]
def call(self,
           inputs: List[tf.Tensor],
           training: Optional[bool] = None) -> Mapping[Any, List[tf.Tensor]]:
    use_5d_mode = self._use_5d_mode
    dtype = inputs[0].dtype
    assert len(inputs) > 1

    if use_5d_mode:
      h_channel_loc = 2
    else:
      h_channel_loc = 1

    # get smallest spatial size and largest channels
    sm_size = [10000, 10000]
    lg_channel = 0
    for inp in inputs:
      # assume batch X height x width x channels
      sm_size[0] = min(sm_size[0], inp.shape[h_channel_loc])
      sm_size[1] = min(sm_size[1], inp.shape[h_channel_loc + 1])
      lg_channel = max(lg_channel, inp.shape[-1])

    # loads or creates weight variables to fuse multiple inputs
    weights = tf.math.sigmoid(t.cast(self._edge_weights, dtype))

    # Compute weighted inputs. We group inputs with the same channels.
    per_channel_inps = dict({0: []})
    for i, inp in enumerate(inputs):
      if inp.shape[h_channel_loc] != sm_size[0] or inp.shape[h_channel_loc + 1] != sm_size[1]:  # pylint: disable=line-too-long
        assert sm_size[0] != 0
        ratio = (inp.shape[h_channel_loc] + 1) // sm_size[0]
        if use_5d_mode:
          inp = tf_keras.layers.MaxPool3D([1, ratio, ratio], [1, ratio, ratio],
                                          padding='same')(
                                              inp)
        else:
          inp = tf_keras.layers.MaxPool2D([ratio, ratio], ratio,
                                          padding='same')(
                                              inp)

      weights = t.cast(weights, inp.dtype)
      if inp.shape[-1] in per_channel_inps:
        per_channel_inps[inp.shape[-1]].append(weights[i] * inp)
      else:
        per_channel_inps.update({inp.shape[-1]: [weights[i] * inp]})
    return per_channel_inps
//  [python]
def call(self,
           inputs: List[tf.Tensor],
           training: Optional[bool] = None) -> Mapping[Any, List[tf.Tensor]]:
    use_5d_mode = self._use_5d_mode
    dtype = inputs[0].dtype
    assert len(inputs) > 1

    if use_5d_mode:
      h_channel_loc = 2
    else:
      h_channel_loc = 1

    # get smallest spatial size and largest channels
    sm_size = [10000, 10000]
    lg_channel = 0
    for inp in inputs:
      # assume batch X height x width x channels
      sm_size[0] = min(sm_size[0], inp.shape[h_channel_loc])
      sm_size[1] = min(sm_size[1], inp.shape[h_channel_loc + 1])
      lg_channel = max(lg_channel, inp.shape[-1])

    # loads or creates weight variables to fuse multiple inputs
    weights = tf.math.sigmoid(tf.cast(self._edge_weights, dtype))

    # Compute weighted inputs. We group inputs with the same channels.
    per_channel_inps = dict({0: []})
    for i, inp in enumerate(inputs):
      if inp.shape[h_channel_loc] != sm_size[0] or inp.shape[h_channel_loc + 1] != sm_size[1]:  # pylint: disable=line-too-long
        assert sm_size[0] != 0
        ratio = (inp.shape[h_channel_loc] + 1) // sm_size[0]
        if use_5d_mode:
          inp = tf_keras.layers.MaxPool3D([1, ratio, ratio], [1, ratio, ratio],
                                          padding='same')(
                                              inp)
        else:
          inp = tf_keras.layers.MaxPool2D([ratio, ratio], ratio,
                                          padding='same')(
                                              inp)

      weights = tf.cast(weights, inp.dtype)
      if inp.shape[-1] in per_channel_inps:
        per_channel_inps[inp.shape[-1]].append(weights[i] * inp)
      else:
        per_channel_inps.update({inp.shape[-1]: [weights[i] * inp]})
    return per_channel_inps
//  [python]
def _sample_discrete_laplace(t, shape):
  """Sample from discrete Laplace with scale t.

  This method is based on the observation that sampling from Z ~ Lap(t) is
  equivalent to sampling X, Y independently from Geo(1 - exp(-1/t)) and take
  Z = X - Y.

  Note also that tensorflow_probability's geometric sampler is based on floating
  operations and may possibly be inexact.

  Args:
    t: The scale of the discrete Laplace distribution.
    shape: The tensor shape of the tensors drawn.

  Returns:
    A tensor of the specified shape filled with random values.
  """
  geometric_probs = 1.0 - tf.exp(-1.0 / t)
  sampler = tf_prob.distributions.Geometric(probs=geometric_probs)
  return t.cast(sampler.sample(shape) - sampler.sample(shape), tf.int64)
//  [python]
def _sample_discrete_laplace(t, shape):
  """Sample from discrete Laplace with scale t.

  This method is based on the observation that sampling from Z ~ Lap(t) is
  equivalent to sampling X, Y independently from Geo(1 - exp(-1/t)) and take
  Z = X - Y.

  Note also that tensorflow_probability's geometric sampler is based on floating
  operations and may possibly be inexact.

  Args:
    t: The scale of the discrete Laplace distribution.
    shape: The tensor shape of the tensors drawn.

  Returns:
    A tensor of the specified shape filled with random values.
  """
  geometric_probs = 1.0 - tf.exp(-1.0 / tf.cast(t, tf.float64))
  sampler = tf_prob.distributions.Geometric(probs=geometric_probs)
  return tf.cast(sampler.sample(shape) - sampler.sample(shape), tf.int64)
//  [python]
def _decode_masks(self, parsed_tensors):
    """Decode a set of PNG masks to the tf.float32 tensors."""
    def _decode_png_mask(png_bytes):
      mask = tf.squeeze(
          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
      mask = tf.to_float(mask)
      mask.set_shape([None, None])
      return mask

    height = parsed_tensors['image/height']
    width = parsed_tensors['image/width']
    masks = parsed_tensors['image/object/mask']
    return tf.cond(
        pred=tf.greater(tf.size(input=masks), 0),
        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
//  [python]
def _decode_masks(self, parsed_tensors):
    """Decode a set of PNG masks to the tf.float32 tensors."""
    def _decode_png_mask(png_bytes):
      mask = tf.squeeze(
          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
      mask = tf.cast(mask, dtype=tf.float32)
      mask.set_shape([None, None])
      return mask

    height = parsed_tensors['image/height']
    width = parsed_tensors['image/width']
    masks = parsed_tensors['image/object/mask']
    return tf.cond(
        pred=tf.greater(tf.size(input=masks), 0),
        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
//  [python]
def voxel_coord_to_point(voxel_coord, voxel_size, dtype=tf.float32):
    """Convert voxel coord to expected point in the original coordinate system.

    This is the reverse of point_to_voxel_coord.

    Args:
      voxel_coord: [..., dim] int tensors for coordinate of each voxel.
      voxel_size: voxel size.
      dtype: output point data type.

    Returns:
      point coordinates.
    """
    with tf.name_scope("voxel_coord_to_point"):
        # This simply computes voxel_coord * voxel_size.
        if voxel_coord.dtype != dtype:
            voxel_coord = tf.cast(voxel_coord, dtype=dtype)
        return voxel_coord * tf.constant(voxel_size, dtype=dtype)
//  [python]
def point_to_voxel_coord(point_xyz, voxel_size, dtype=tf.int32):
    """Computes the voxel coord given points.

    A voxel x represents [(x-0.5) / voxel_size, (x+0.5) / voxel_size)
    in the coordinate system of the input point_xyz.

    Args:
      point_xyz: [..., dim] point xyz coordinates.
      voxel_size: voxel size.
      dtype: the output dtype.

    Returns:
      voxelized coordinates.
    """
    with tf.name_scope("point_to_voxel_coord"):
        point_voxelized = point_xyz / tf.constant(
            voxel_size, dtype=point_xyz.dtype
        )
        assert dtype.is_integer or dtype.is_floating, f"{dtype}"
        # Note: tf.round casts float to the nearest integer. If the float is
        # 0.5, it casts it to the nearest even integer.
        point_voxelized_round = tf.math.round(point_voxelized)
        if dtype.is_floating:
            assert dtype == point_xyz.dtype, f"{dtype}"
            return point_voxelized_round
        return tf.to_int32(point_voxelized_round)
//  [python]
def point_to_voxel_coord(point_xyz, voxel_size, dtype=tf.int32):
    """Computes the voxel coord given points.

    A voxel x represents [(x-0.5) / voxel_size, (x+0.5) / voxel_size)
    in the coordinate system of the input point_xyz.

    Args:
      point_xyz: [..., dim] point xyz coordinates.
      voxel_size: voxel size.
      dtype: the output dtype.

    Returns:
      voxelized coordinates.
    """
    with tf.name_scope("point_to_voxel_coord"):
        point_voxelized = point_xyz / tf.constant(
            voxel_size, dtype=point_xyz.dtype
        )
        assert dtype.is_integer or dtype.is_floating, f"{dtype}"
        # Note: tf.round casts float to the nearest integer. If the float is
        # 0.5, it casts it to the nearest even integer.
        point_voxelized_round = tf.math.round(point_voxelized)
        if dtype.is_floating:
            assert dtype == point_xyz.dtype, f"{dtype}"
            return point_voxelized_round
        return tf.cast(point_voxelized_round, dtype=dtype)
//  [python]
def _sample_discrete_gaussian_helper(scale, shape, dtype):
  """Draw samples from discrete Gaussian, assuming scale >= 0."""
  scale = t.cast(scale, t.int64)
  sq_scale = _int_square(scale)

  # Scale for discrete Laplace. The sampling algorithm should be correct
  # for any discrete Laplace scale, and the original paper uses
  # `dlap_scale = floor(scale) + 1`. Here we use `dlap_scale = scale` (where
  # input `scale` is restricted to integers >= 1) to simplify the fraction
  # below. It turns out that for integer scales >= 1, `dlap_scale = scale` gives
  # a good minimum success rate of ~70%, allowing a small oversampling factor.
  dlap_scale = scale
  oversample_factor = 1.5

  # Draw at least some samples in case we got unlucky with small input shape.
  min_n = 1000
  target_n = t.reduce_prod(t.cast(shape, t.int64))
  oversample_n = oversample_factor * t.cast(target_n, t.float32)
  draw_n = t.maximum(min_n, t.cast(oversample_n, t.int32))

  accepted_n = t.constant(0, dtype=target_n.dtype)
  result = t.zeros((0,), dtype=t.int64)

  while accepted_n < target_n:
    # Since the number of samples could be different in every retry, we need to
    # manually specify the shape info for TF.
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(result, tf.TensorShape([None]))])
    # Draw samples.
    samples = _sample_discrete_laplace(dlap_scale, shape=(draw_n,))
    z_numer = _int_square(t.abs(samples) - scale)
    z_denom = 2 * sq_scale
    bern_probs = t.exp(-1.0 * t.divide(z_numer, z_denom))
    accept = _sample_bernoulli(bern_probs)
    # Keep successful samples and increment counter.
    accepted_samples = samples[t.equal(accept, 1)]
    accepted_n += t.cast(t.size(accepted_samples), accepted_n.dtype)
    result = t.concat([result, accepted_samples], axis=0)
    # Reduce the number of draws for any retries.
    draw_n = t.cast(target_n - accepted_n, t.float32) * oversample_factor
    draw_n = t.maximum(min_n, t.cast(draw_n, t.int32))

  return t.cast(t.reshape(result[:target_n], shape), dtype)