jiahuimbzuai/CoQuIR huggingface dataset:
//  [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:
        ----
            gt: the encoded tensor with gt labels
            model_output: predicted logits of the model
            seq_len: lengths of each gt word inside the batch

        Returns:
        -------
            The loss of the model on the batch
        """
        # Input length : number of timesteps
        input_len = tf.shape(model_output)[1]
        # Add one for additional <eos> token (sos disappear in shift!)
        seq_len = t.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>. Delete last logit of the model output.
        cce = tf.nn.softmax_cross_entropy_with_logits(oh_gt[:, 1:, :], model_output[:, :-1, :])
        # Compute mask
        mask_values = tf.zeros_like(cce)
        mask_2d = tf.sequence_mask(seq_len, input_len - 1)  # delete the last mask timestep as well
        masked_loss = tf.where(mask_2d, cce, mask_values)
        ce_loss = tf.math.divide(tf.reduce_sum(masked_loss, axis=1), t.cast(seq_len, model_output.dtype))

        return tf.expand_dims(ce_loss, axis=1)
//  [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:
        ----
            gt: the encoded tensor with gt labels
            model_output: predicted logits of the model
            seq_len: lengths of each gt word inside the batch

        Returns:
        -------
            The loss of the model on the batch
        """
        # Input length : number of timesteps
        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>. Delete last logit of the model output.
        cce = tf.nn.softmax_cross_entropy_with_logits(oh_gt[:, 1:, :], model_output[:, :-1, :])
        # Compute mask
        mask_values = tf.zeros_like(cce)
        mask_2d = tf.sequence_mask(seq_len, input_len - 1)  # delete the last mask timestep as well
        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 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.
    """
    target_class_ids = tf.to_int64(target_class_ids)
    #target_class_ids = tf.Print(target_class_ids, [target_class_ids], message="target_class_ids", summarize=15)

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(pred_class_logits, axis=2)
    # TODO: Update this line to work with batch > 1. Right now it assumes all
    #       images in a batch have the same active_class_ids
    #pred_class_ids = tf.Print(pred_class_ids, [pred_class_logits], message="pred_class_logits", summarize=200)
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)
    #pred_active = tf.Print(pred_active, [pred_active], message="pred_active", summarize=15)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)
    #loss = tf.Print(loss, [loss], message="loss", summarize=15)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active
    #loss = tf.Print(loss, [loss], message="loss_after_multiply", summarize=15)

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    #loss = tf.Print(loss, [loss, tf.reduce_sum(pred_active), tf.reduce_sum(loss)], message="loss_after_mean",
    #                summarize=15)
    loss = K.switch(tf.reduce_sum(pred_active) > 0, tf.reduce_sum(loss)/ tf.reduce_sum(pred_active), tf.constant(0.0))

    ##tf.reduce_sum(loss) / tf.reduce_sum(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.
    """
    target_class_ids = tf.cast(target_class_ids, 'int64')
    #target_class_ids = tf.Print(target_class_ids, [target_class_ids], message="target_class_ids", summarize=15)

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(pred_class_logits, axis=2)
    # TODO: Update this line to work with batch > 1. Right now it assumes all
    #       images in a batch have the same active_class_ids
    #pred_class_ids = tf.Print(pred_class_ids, [pred_class_logits], message="pred_class_logits", summarize=200)
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)
    #pred_active = tf.Print(pred_active, [pred_active], message="pred_active", summarize=15)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)
    #loss = tf.Print(loss, [loss], message="loss", summarize=15)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active
    #loss = tf.Print(loss, [loss], message="loss_after_multiply", summarize=15)

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    #loss = tf.Print(loss, [loss, tf.reduce_sum(pred_active), tf.reduce_sum(loss)], message="loss_after_mean",
    #                summarize=15)
    loss = K.switch(tf.reduce_sum(pred_active) > 0, tf.reduce_sum(loss)/ tf.reduce_sum(pred_active), tf.constant(0.0))

    ##tf.reduce_sum(loss) / tf.reduce_sum(pred_active)

    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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(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(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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]

    binary_crossentropy_loss = K.binary_crossentropy(target=y_true, output=y_pred)
    loss = K.switch(tf.size(y_true) > 0,
                    binary_crossentropy_loss,
                    tf.constant(0.0))

    loss = K.mean(loss)
    loss = K.reshape(loss, [1, 1])
    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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(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(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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]

    binary_crossentropy_loss = K.binary_crossentropy(target=y_true, output=y_pred)
    loss = K.switch(tf.size(y_true) > 0,
                    binary_crossentropy_loss,
                    tf.constant(0.0))

    loss = K.mean(loss)
    loss = K.reshape(loss, [1, 1])
    return loss
//  [python]
def mrcnn_coord_delta_index(mrcnn_coord_delta, mrcnn_coord_bin):
    """
        Might be for x, y or z channel.
        mrcnn_coord_delta: [batch, proposals, height, width, num_classes, num_bins].
        mrcnn_coord_bin: [batch, proposals, height, width, num_classes, 1].
    """

    shape = tf.shape(mrcnn_coord_delta)
    reshape_params = tf.reshape(mrcnn_coord_delta, [-1, tf.shape(mrcnn_coord_delta)[-1]])

    reshape_indices = tf.to_int64(tf.reshape(mrcnn_coord_bin, [-1]))

    nums = tf.to_int64(tf.range(tf.shape(reshape_params)[0]))
    new_indice = tf.stack([nums, reshape_indices], axis=1)

    # X, Y = tf.meshgrid(num_0, num_1)
    # new_indices = tf.stack([Y, X, indices], axis=2)


    output = tf.gather_nd(reshape_params, new_indice)
    output = tf.reshape(output, shape[:-1])

    return output
//  [python]
def mrcnn_coord_delta_index(mrcnn_coord_delta, mrcnn_coord_bin):
    """
        Might be for x, y or z channel.
        mrcnn_coord_delta: [batch, proposals, height, width, num_classes, num_bins].
        mrcnn_coord_bin: [batch, proposals, height, width, num_classes, 1].
    """

    shape = tf.shape(mrcnn_coord_delta)
    reshape_params = tf.reshape(mrcnn_coord_delta, [-1, tf.shape(mrcnn_coord_delta)[-1]])

    reshape_indices = tf.cast(tf.reshape(mrcnn_coord_bin, [-1]), tf.int64)

    nums = tf.cast(tf.range(tf.shape(reshape_params)[0]), tf.int64)
    new_indice = tf.stack([nums, reshape_indices], axis=1)

    # X, Y = tf.meshgrid(num_0, num_1)
    # new_indices = tf.stack([Y, X, indices], axis=2)


    output = tf.gather_nd(reshape_params, new_indice)
    output = tf.reshape(output, shape[:-1])

    return output
//  [python]
def trim_zeros_graph(boxes):
    """Often boxes are represented with matricies of shape [N, 4] and
    are padded with zeros. This removes zero boxes.

    boxes: [N, 4] matrix of boxes.

    TODO: use this function to reduce code duplication
    """
    area = tf.boolean_mask(boxes, tf.cast(tf.reduce_sum(tf.abs(boxes), axis=1),
                           tf.bool))
//  [python]
def mrcnn_coord_l1_loss_graph(target_masks, target_coord, target_class_ids, pred_coord):
    """Mask L1 loss for the coordinates head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_coord: [batch, num_rois, height, width]. Might be for x, y or z channel.
        A float32 tensor of values in the range of [0, 1]. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_coord: [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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    target_coord = K.reshape(target_coord, (-1, mask_shape[2], mask_shape[3]))

    pred_shape = tf.shape(pred_coord)
    pred_coord = K.reshape(pred_coord, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_coord = tf.transpose(pred_coord, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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 coordinate maps and masks (predicted and true) that contribute to loss
    # true coord map:[N', height, width]
    y_true = tf.gather(target_coord, positive_ix) ## shape: [num_pos_rois, height, width]
    mask = tf.gather(target_masks, positive_ix) ## shape: [num_pos_rois, height, width]
    mask = tf.cast(mask, dtype=tf.bool)
    y_true_in_mask = tf.boolean_mask(y_true, mask) ## shape: [num_pos_rois, height, width]


    #y_mask = tf.gather(target_masks, positive_ix)
    #num_of_pixels = tf.reduce_sum(y_mask, axis=[1, 2]) + 0.00001
    #y_pred = tf.gather_nd(pred_coord, indices)
    #y_pred_in_mask = tf.multiply(y_mask, y_pred)

    y_pred = tf.gather_nd(pred_coord, indices)
    y_pred_in_mask = tf.boolean_mask(y_pred, mask)

    #coord_loss = K.sum(K.abs(y_true_in_mask - y_pred_in_mask), axis=[1, 2])

    coord_loss = K.abs(y_true_in_mask - y_pred_in_mask)
    mean_loss =  K.mean(coord_loss)

    loss = K.switch(tf.size(y_true) > 0, mean_loss, tf.constant(0.0))
    loss = K.reshape(loss, [1, 1])

    return loss
//  [python]
def mrcnn_coord_l1_loss_graph(target_masks, target_coord, target_class_ids, pred_coord):
    """Mask L1 loss for the coordinates head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_coord: [batch, num_rois, height, width]. Might be for x, y or z channel.
        A float32 tensor of values in the range of [0, 1]. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_coord: [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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    target_coord = K.reshape(target_coord, (-1, mask_shape[2], mask_shape[3]))

    pred_shape = tf.shape(pred_coord)
    pred_coord = K.reshape(pred_coord, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_coord = tf.transpose(pred_coord, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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 coordinate maps and masks (predicted and true) that contribute to loss
    # true coord map:[N', height, width]
    y_true = tf.gather(target_coord, positive_ix) ## shape: [num_pos_rois, height, width]
    mask = tf.gather(target_masks, positive_ix) ## shape: [num_pos_rois, height, width]
    mask = tf.cast(mask, dtype=tf.bool)
    y_true_in_mask = tf.boolean_mask(y_true, mask) ## shape: [num_pos_rois, height, width]


    #y_mask = tf.gather(target_masks, positive_ix)
    #num_of_pixels = tf.reduce_sum(y_mask, axis=[1, 2]) + 0.00001
    #y_pred = tf.gather_nd(pred_coord, indices)
    #y_pred_in_mask = tf.multiply(y_mask, y_pred)

    y_pred = tf.gather_nd(pred_coord, indices)
    y_pred_in_mask = tf.boolean_mask(y_pred, mask)

    #coord_loss = K.sum(K.abs(y_true_in_mask - y_pred_in_mask), axis=[1, 2])

    coord_loss = K.abs(y_true_in_mask - y_pred_in_mask)
    mean_loss =  K.mean(coord_loss)

    loss = K.switch(tf.size(y_true) > 0, mean_loss, tf.constant(0.0))
    loss = K.reshape(loss, [1, 1])

    return loss
//  [python]
def reweight(msa1hot, cutoff):

    with tf.name_scope('reweight'):
        id_min = tf.cast(tf.shape(msa1hot)[1], tf.float32) * cutoff
        id_mtx = tf.tensordot(msa1hot, msa1hot, [[1,2], [1,2]])
        id_mask = id_mtx > id_min
        w = 1.0/tf.reduce_sum(tf.cast(id_mask, dtype=tf.float32),-1)
    return w
//  [python]
def mrcnn_coord_smooth_l1_loss_graph(target_masks, target_coord, target_class_ids, pred_coord):
    """Mask L1 loss for the coordinates head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_coord: [batch, num_rois, height, width]. Might be for x, y or z channel.
        A float32 tensor of values in the range of [0, 1]. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_coord: [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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    target_coord = K.reshape(target_coord, (-1, mask_shape[2], mask_shape[3]))

    pred_shape = tf.shape(pred_coord)
    pred_coord = K.reshape(pred_coord, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_coord = tf.transpose(pred_coord, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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 coordinate maps and masks (predicted and true) that contribute to loss
    # true coord map:[N', height, width]
    y_true = tf.gather(target_coord, positive_ix)  ## shape: [num_pos_rois, height, width]
    mask = tf.gather(target_masks, positive_ix)  ## shape: [num_pos_rois, height, width]
    mask = tf.cast(mask, dtype=tf.bool)
    y_true_in_mask = tf.boolean_mask(y_true, mask)  ## shape: [num_pos_rois, height, width]

    # y_mask = tf.gather(target_masks, positive_ix)
    # num_of_pixels = tf.reduce_sum(y_mask, axis=[1, 2]) + 0.00001
    # y_pred = tf.gather_nd(pred_coord, indices)
    # y_pred_in_mask = tf.multiply(y_mask, y_pred)

    y_pred = tf.gather_nd(pred_coord, indices)
    y_pred_in_mask = tf.boolean_mask(y_pred, mask)

    # coord_loss = K.sum(K.abs(y_true_in_mask - y_pred_in_mask), axis=[1, 2])

    diff = K.abs(y_true_in_mask - y_pred_in_mask)
    threshold = 0.1
    coefficient = 1/(2*threshold)
    less_than_threshold = K.cast(K.less(diff, threshold), "float32")
    loss = (less_than_threshold * coefficient * diff**2) + (1-less_than_threshold) * (diff - threshold/2)

    return loss
//  [python]
def mrcnn_coord_reg_loss_graph(target_masks, target_coord, target_class_ids, pred_coord, loss_fn):
    """Mask L1 loss for the coordinates head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_coord: [batch, num_rois, height, width]. Might be for x, y or z channel.
        A float32 tensor of values in the range of [0, 1]. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_coord: [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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    target_coord = K.reshape(target_coord, (-1, mask_shape[2], mask_shape[3]))

    pred_shape = tf.shape(pred_coord)
    pred_coord = K.reshape(pred_coord, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_coord = tf.transpose(pred_coord, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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 coordinate maps and masks (predicted and true) that contribute to loss
    # true coord map:[N', height, width]
    y_true = tf.gather(target_coord, positive_ix)  ## shape: [num_pos_rois, height, width]
    mask = tf.gather(target_masks, positive_ix)  ## shape: [num_pos_rois, height, width]
    mask = tf.cast(mask, dtype=tf.bool)

    #assert_op = tf.Assert(tf.greater(tf.reduce_max(tf.cast(mask, dtype=tf.float32)), 0.),
    #                      [tf.size(mask), tf.size(positive_ix), target_class_ids])

    #with tf.control_dependencies([assert_op]):

    y_true_in_mask = tf.boolean_mask(y_true, mask)  ## shape: [num_pixels_in_masks_for_all_pos_rois]

    # y_mask = tf.gather(target_masks, positive_ix)
    # num_of_pixels = tf.reduce_sum(y_mask, axis=[1, 2]) + 0.00001
    # y_pred = tf.gather_nd(pred_coord, indices)
    # y_pred_in_mask = tf.multiply(y_mask, y_pred)

    y_pred = tf.gather_nd(pred_coord, indices)
    y_pred_in_mask = tf.boolean_mask(y_pred, mask)

    # coord_loss = K.sum(K.abs(y_true_in_mask - y_pred_in_mask), axis=[1, 2])

    diff = K.abs(y_true_in_mask - y_pred_in_mask)
    loss = loss_fn(diff)

    loss = K.switch(tf.size(loss) > 0, K.mean(loss), tf.constant(0.0))

    return loss
//  [python]
def mrcnn_coord_reg_loss_graph(target_masks, target_coord, target_class_ids, pred_coord, loss_fn):
    """Mask L1 loss for the coordinates head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_coord: [batch, num_rois, height, width]. Might be for x, y or z channel.
        A float32 tensor of values in the range of [0, 1]. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_coord: [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(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    target_coord = K.reshape(target_coord, (-1, mask_shape[2], mask_shape[3]))

    pred_shape = tf.shape(pred_coord)
    pred_coord = K.reshape(pred_coord, (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_coord = tf.transpose(pred_coord, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.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 coordinate maps and masks (predicted and true) that contribute to loss
    # true coord map:[N', height, width]
    y_true = tf.gather(target_coord, positive_ix)  ## shape: [num_pos_rois, height, width]
    mask = tf.gather(target_masks, positive_ix)  ## shape: [num_pos_rois, height, width]
    mask = tf.cast(mask, dtype=tf.bool)

    #assert_op = tf.Assert(tf.greater(tf.reduce_max(tf.cast(mask, dtype=tf.float32)), 0.),
    #                      [tf.size(mask), tf.size(positive_ix), target_class_ids])

    #with tf.control_dependencies([assert_op]):

    y_true_in_mask = tf.boolean_mask(y_true, mask)  ## shape: [num_pixels_in_masks_for_all_pos_rois]

    # y_mask = tf.gather(target_masks, positive_ix)
    # num_of_pixels = tf.reduce_sum(y_mask, axis=[1, 2]) + 0.00001
    # y_pred = tf.gather_nd(pred_coord, indices)
    # y_pred_in_mask = tf.multiply(y_mask, y_pred)

    y_pred = tf.gather_nd(pred_coord, indices)
    y_pred_in_mask = tf.boolean_mask(y_pred, mask)

    # coord_loss = K.sum(K.abs(y_true_in_mask - y_pred_in_mask), axis=[1, 2])

    diff = K.abs(y_true_in_mask - y_pred_in_mask)
    loss = loss_fn(diff)

    loss = K.switch(tf.size(loss) > 0, K.mean(loss), tf.constant(0.0))

    return loss
//  [python]
def read_tensor_from_image_file(self, file_name, input_height=299, input_width=299, input_mean=0, input_std=255):
        input_name = "file_reader"
        try:

            file_reader = tf.io.read_file(file_name, input_name)
            if file_name.endswith(".png"):
                image_reader = tf.image.decode_png(file_reader, channels=3, name='png_reader')
            elif file_name.endswith(".gif"):
                image_reader = tf.squeeze(tf.image.decode_gif(file_reader, name='gif_reader'))
            elif file_name.endswith(".bmp"):
                image_reader = tf.image.decode_bmp(file_reader, name='bmp_reader')
            else:
                image_reader = tf.image.decode_jpeg(file_reader, channels=3, name='jpeg_reader')
            float_caster = tf.to_float(image_reader)
            dims_expander = tf.expand_dims(float_caster, 0)
            resized = tf.image.resize(dims_expander, [input_height, input_width], method=tf.image.ResizeMethod.BILINEAR, antialias=True)
            normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
            sess = tf.compat.v1.Session()
            return sess.run(normalized)
        except:
            return None
//  [python]
def read_tensor_from_image_file(self, file_name, input_height=299, input_width=299, input_mean=0, input_std=255):
        input_name = "file_reader"
        try:

            file_reader = tf.io.read_file(file_name, input_name)
            if file_name.endswith(".png"):
                image_reader = tf.image.decode_png(file_reader, channels=3, name='png_reader')
            elif file_name.endswith(".gif"):
                image_reader = tf.squeeze(tf.image.decode_gif(file_reader, name='gif_reader'))
            elif file_name.endswith(".bmp"):
                image_reader = tf.image.decode_bmp(file_reader, name='bmp_reader')
            else:
                image_reader = tf.image.decode_jpeg(file_reader, channels=3, name='jpeg_reader')
            float_caster = tf.cast(image_reader, tf.float32)
            dims_expander = tf.expand_dims(float_caster, 0)
            resized = tf.image.resize(dims_expander, [input_height, input_width], method=tf.image.ResizeMethod.BILINEAR, antialias=True)
            normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
            sess = tf.compat.v1.Session()
            return sess.run(normalized)
        except:
            return None
//  [python]
def get_features1d(msa1hot,insertions,w):
    nc = tf.shape(msa1hot)[1]
    beff = tf.reduce_sum(w)
    msw = w[:,None,None]*msa1hot
    f_i = tf.reduce_sum(msw, axis=0) / beff + 1e-9
    h_i = tf.reduce_sum( -f_i * tf.math.log(f_i), axis=1)
    n_i = tf.math.log(tf.reduce_sum(msw[:,:,:20],axis=[0,2]))
    ins = tf.reduce_mean(tf.to_float(insertions>0),axis=0)
    f1d = tf.concat([msa1hot[0,:,:20], f_i, h_i[:,None], n_i[:,None], ins[:,None]], axis=1)
    f1d = tf.reshape(f1d, [1,nc,44])
    return f1d
//  [python]
def get_features1d(msa1hot,insertions,w):
    nc = tf.shape(msa1hot)[1]
    beff = tf.reduce_sum(w)
    msw = w[:,None,None]*msa1hot
    f_i = tf.reduce_sum(msw, axis=0) / beff + 1e-9
    h_i = tf.reduce_sum( -f_i * tf.math.log(f_i), axis=1)
    n_i = tf.math.log(tf.reduce_sum(msw[:,:,:20],axis=[0,2]))
    ins = tf.reduce_mean(tf.cast(insertions>0,dtype=tf.float32),axis=0)
    f1d = tf.concat([msa1hot[0,:,:20], f_i, h_i[:,None], n_i[:,None], ins[:,None]], axis=1)
    f1d = tf.reshape(f1d, [1,nc,44])
    return f1d
//  [python]
def get_features(msa, ins, idx, params):

    nrow = tf.shape(msa)[0]
    ncol = tf.shape(msa)[1]

    # convert msa to 1-hot
    msa1hot  = tf.one_hot(msa, 21, dtype=tf.float32)

    # extract 1d and 2d features from the msa
    w = reweight(msa1hot, params['WMIN'])
    f1d = get_features1d(msa1hot,ins,w)
    f2d = tf.cond(nrow>1, lambda: get_features2d(msa1hot, w, params['DCAREG']),
            lambda: tf.zeros([1,ncol,ncol,444], tf.float32))

    # get sequence separation
    seqsep = tf.abs(idx[:,None]-idx[None,:])+1
    seqsep = tf.math.log(tf.cast(seqsep,dtype=tf.float32))

    # stack all features together
    inputs = tf.concat([tf.tile(f1d[:,:,None,:], [1,1,ncol,1]),
                        tf.tile(f1d[:,None,:,:], [1,ncol,1,1]),
                        f2d,
                        seqsep[None,:,:,None]], axis=-1)
    inputs = tf.reshape(inputs, [1,ncol,ncol,2*44+444+1])

    return inputs
//  [python]
def _batch_count_ones(masks: tf.Tensor,
                      dtype: tf.dtypes.DType = tf.int32) -> tf.Tensor:
  """Counts the ones/trues for each mask in the batch.

  Args:
    masks: A tensor in shape (..., height, width) with arbitrary numbers of
      batch dimensions.
    dtype: DType of the resulting tensor. Default is tf.int32.

  Returns:
    A tensor which contains the count of non-zero elements for each mask in the
    batch. The rank of the resulting tensor is equal to rank(masks) - 2.
  """
  masks_shape = masks.get_shape().as_list()
  if len(masks_shape) < 2:
    raise ValueError(
        'Expected the input masks (..., height, width) has rank >= 2, was: %s' %
        masks_shape)
  return tf.reduce_sum(tf.to_int32(masks), axis=[-2, -1])
//  [python]
def _batch_count_ones(masks: tf.Tensor,
                      dtype: tf.dtypes.DType = tf.int32) -> tf.Tensor:
  """Counts the ones/trues for each mask in the batch.

  Args:
    masks: A tensor in shape (..., height, width) with arbitrary numbers of
      batch dimensions.
    dtype: DType of the resulting tensor. Default is tf.int32.

  Returns:
    A tensor which contains the count of non-zero elements for each mask in the
    batch. The rank of the resulting tensor is equal to rank(masks) - 2.
  """
  masks_shape = masks.get_shape().as_list()
  if len(masks_shape) < 2:
    raise ValueError(
        'Expected the input masks (..., height, width) has rank >= 2, was: %s' %
        masks_shape)
  return tf.reduce_sum(tf.cast(masks, dtype), axis=[-2, -1])
//  [python]
def call(self, y_true, y_pred):
    """Invokes the `FocalLoss`.

    Args:
      y_true: Ordered Dict with level to [batch, height, width, num_anchors].
        for example,
        {3: tf.Tensor(shape=[32, 512, 512, 9], dtype=tf.float32),
         4: tf.Tensor([shape=32, 256, 256, 9, dtype=tf.float32])}
      y_pred: Ordered Dict with level to [batch, height, width, num_anchors *
        num_classes]. for example,
        {3: tf.Tensor(shape=[32, 512, 512, 9], dtype=tf.int64),
         4: tf.Tensor(shape=[32, 256, 256, 9 * 21], dtype=tf.int64)}

    Returns:
      Summed loss float `Tensor`.
    """
    flattened_cls_outputs = []
    flattened_labels = []
    batch_size = None
    for level in y_pred.keys():
      cls_output = y_pred[level]
      label = y_true[level]
      if batch_size is None:
        batch_size = cls_output.shape[0] or tf.shape(cls_output)[0]
      flattened_cls_outputs.append(
          tf.reshape(cls_output, [batch_size, -1, self._num_classes]))
      flattened_labels.append(tf.reshape(label, [batch_size, -1]))
    cls_outputs = tf.concat(flattened_cls_outputs, axis=1)
    labels = tf.concat(flattened_labels, axis=1)

    cls_targets_one_hot = tf.one_hot(labels, self._num_classes)
    return focal_loss(
        tf.to_float(cls_outputs),
        tf.to_float(cls_targets_one_hot), self._alpha,
        self._gamma)
//  [python]
def call(self, y_true, y_pred):
    """Invokes the `FocalLoss`.

    Args:
      y_true: Ordered Dict with level to [batch, height, width, num_anchors].
        for example,
        {3: tf.Tensor(shape=[32, 512, 512, 9], dtype=tf.float32),
         4: tf.Tensor([shape=32, 256, 256, 9, dtype=tf.float32])}
      y_pred: Ordered Dict with level to [batch, height, width, num_anchors *
        num_classes]. for example,
        {3: tf.Tensor(shape=[32, 512, 512, 9], dtype=tf.int64),
         4: tf.Tensor(shape=[32, 256, 256, 9 * 21], dtype=tf.int64)}

    Returns:
      Summed loss float `Tensor`.
    """
    flattened_cls_outputs = []
    flattened_labels = []
    batch_size = None
    for level in y_pred.keys():
      cls_output = y_pred[level]
      label = y_true[level]
      if batch_size is None:
        batch_size = cls_output.shape[0] or tf.shape(cls_output)[0]
      flattened_cls_outputs.append(
          tf.reshape(cls_output, [batch_size, -1, self._num_classes]))
      flattened_labels.append(tf.reshape(label, [batch_size, -1]))
    cls_outputs = tf.concat(flattened_cls_outputs, axis=1)
    labels = tf.concat(flattened_labels, axis=1)

    cls_targets_one_hot = tf.one_hot(labels, self._num_classes)
    return focal_loss(
        tf.cast(cls_outputs, dtype=tf.float32),
        tf.cast(cls_targets_one_hot, dtype=tf.float32), self._alpha,
        self._gamma)
//  [python]
def permutation_loss(pred_dsmat: tf.Tensor, gt_perm: tf.Tensor, n1: tf.Tensor, n2: tf.Tensor) -> tf.Tensor:
    """
    Tensorflow implementation of permutation_loss
    """
    batch_num = pred_dsmat.shape[0]

    pred_dsmat = t.cast(pred_dsmat, dtype=tf.float32)

    if not tf.reduce_all(tf.math.logical_and((pred_dsmat >= 0), (pred_dsmat <= 1))):
        raise ValueError("pred_dsmat contains invalid numerical entries.")
    if not tf.reduce_all(tf.math.logical_and((gt_perm >= 0), (gt_perm <= 1))):
        raise ValueError("gt_perm contains invalid numerical entries.")

    if n1 is None:
        n1 = tf.constant([pred_dsmat.shape[1] for _ in range(batch_num)])
    if n2 is None:
        n2 = tf.constant([pred_dsmat.shape[2] for _ in range(batch_num)])

    loss = tf.Variable(0.)
    n_sum = tf.zeros_like(loss)
    for b in range(batch_num):
        batch_slice = [b, slice(n1[b]), slice(n2[b])]
        bce = tf.losses.BinaryCrossentropy()
        loss.assign_add(t.cast(tf.size(gt_perm[batch_slice]), dtype=tf.float32) * bce(
            gt_perm[batch_slice],
            pred_dsmat[batch_slice]))
        n_sum += t.cast(n1[b], dtype=n_sum.dtype)

    return loss / n_sum
//  [python]
def permutation_loss(pred_dsmat: tf.Tensor, gt_perm: tf.Tensor, n1: tf.Tensor, n2: tf.Tensor) -> tf.Tensor:
    """
    Tensorflow implementation of permutation_loss
    """
    batch_num = pred_dsmat.shape[0]

    pred_dsmat = tf.cast(pred_dsmat, dtype=tf.float32)

    if not tf.reduce_all(tf.math.logical_and((pred_dsmat >= 0), (pred_dsmat <= 1))):
        raise ValueError("pred_dsmat contains invalid numerical entries.")
    if not tf.reduce_all(tf.math.logical_and((gt_perm >= 0), (gt_perm <= 1))):
        raise ValueError("gt_perm contains invalid numerical entries.")

    if n1 is None:
        n1 = tf.constant([pred_dsmat.shape[1] for _ in range(batch_num)])
    if n2 is None:
        n2 = tf.constant([pred_dsmat.shape[2] for _ in range(batch_num)])

    loss = tf.Variable(0.)
    n_sum = tf.zeros_like(loss)
    for b in range(batch_num):
        batch_slice = [b, slice(n1[b]), slice(n2[b])]
        bce = tf.losses.BinaryCrossentropy()
        loss.assign_add(tf.cast(tf.size(gt_perm[batch_slice]), dtype=tf.float32) * bce(
            gt_perm[batch_slice],
            pred_dsmat[batch_slice]))
        n_sum += tf.cast(n1[b], dtype=n_sum.dtype)

    return loss / n_sum
//  [python]
def _check_and_init_gm(K, n1, n2, n1max, n2max, x0):
    # get batch number
    batch_num = K.shape[0]
    n1n2 = K.shape[1]

    # get values of n1, n2, n1max, n2max and check
    if n1 is None:
        n1 = t.cast(t.fill((batch_num,), n1max), dtype=t.int32)
    if n2 is None:
        n2 = t.cast(t.fill((batch_num,), n2max), dtype=t.int32)
    if n1max is None:
        n1max = t.reduce_max(n1)
    if n2max is None:
        n2max = t.reduce_max(n2)

    if not n1max * n2max == n1n2:
        raise ValueError('the input size of K does not match with n1max * n2max!')

    # initialize x0 (also v0)
    if x0 is None:
        with t.device(K.device):
            x0 = t.Variable(t.cast(t.zeros([batch_num, n1max, n2max]), dtype=K.dtype))
        for b in range(batch_num):
            f = t.cast(t.fill([n1[b], n2[b]], 1) / (n1[b] * n2[b]), dtype=t.float32)
            x0[b, 0:n1[b], 0:n2[b]].assign(f)
    v0 = t.reshape(t.transpose(x0, perm=[0, 2, 1]), [batch_num, n1n2, 1])

    return batch_num, n1, n2, n1max, n2max, n1n2, v0
//  [python]
def _check_and_init_gm(K, n1, n2, n1max, n2max, x0):
    # get batch number
    batch_num = K.shape[0]
    n1n2 = K.shape[1]

    # get values of n1, n2, n1max, n2max and check
    if n1 is None:
        n1 = tf.cast(tf.fill((batch_num,), n1max), dtype=tf.int32)
    if n2 is None:
        n2 = tf.cast(tf.fill((batch_num,), n2max), dtype=tf.int32)
    if n1max is None:
        n1max = tf.reduce_max(n1)
    if n2max is None:
        n2max = tf.reduce_max(n2)

    if not n1max * n2max == n1n2:
        raise ValueError('the input size of K does not match with n1max * n2max!')

    # initialize x0 (also v0)
    if x0 is None:
        with tf.device(K.device):
            x0 = tf.Variable(tf.cast(tf.zeros([batch_num, n1max, n2max]), dtype=K.dtype))
        for b in range(batch_num):
            f = tf.cast(tf.fill([n1[b], n2[b]], 1) / (n1[b] * n2[b]), dtype=tf.float32)
            x0[b, 0:n1[b], 0:n2[b]].assign(f)
    v0 = tf.reshape(tf.transpose(x0, perm=[0, 2, 1]), [batch_num, n1n2, 1])

    return batch_num, n1, n2, n1max, n2max, n1n2, v0
//  [python]
def contrastive_loss(model1, model2, y, margin):
    with tf.name_scope("contrastive-loss"):
        y = tf.cast(y,tf.float32)
        d = tf.sqrt(tf.reduce_sum(tf.pow(model1-model2, 2), 1, keep_dims=False))
        tmp= y * tf.square(d)
        tmp2 = (1 - y) * tf.square(tf.maximum((margin - d),0))
    return tf.reduce_mean(tmp + tmp2)/2
//  [python]
def siamese_loss(out1,out2,y,Q=5):
    y = t.cast(y,tf.float32)
    Q = tf.constant(Q, name="Q",dtype=tf.float32)
    E_w = tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(out1,out2)),1))   
    pos = tf.multiply(tf.multiply(y,2/Q),tf.square(E_w))
    neg = tf.multiply(tf.multiply(1-y,2*Q),tf.exp(-2.77/Q*E_w))                
    loss = pos + neg                 
    loss = tf.reduce_mean(loss)              
    return loss
//  [python]
def siamese_loss(out1,out2,y,Q=5):
    y = tf.cast(y,tf.float32)
    Q = tf.constant(Q, name="Q",dtype=tf.float32)
    E_w = tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(out1,out2)),1))   
    pos = tf.multiply(tf.multiply(y,2/Q),tf.square(E_w))
    neg = tf.multiply(tf.multiply(1-y,2*Q),tf.exp(-2.77/Q*E_w))                
    loss = pos + neg                 
    loss = tf.reduce_mean(loss)              
    return loss
//  [python]
def call(self,
           inputs: tf.Tensor,
           image_info: Optional[tf.Tensor] = None) -> Dict[str, tf.Tensor]:
    detections = inputs

    batched_scores = detections['detection_scores']
    batched_classes = detections['detection_classes']
    batched_detections_masks = tf.expand_dims(
        detections['detection_masks'], axis=-1)
    batched_boxes = detections['detection_boxes']
    batched_segmentation_masks = tf.to_float(
        detections['segmentation_outputs'])

    if self._rescale_predictions:
      scale = tf.tile(
          tf.cast(image_info[:, 2:3, :], dtype=batched_boxes.dtype),
          multiples=[1, 1, 2])
      batched_boxes /= scale

      batched_segmentation_masks = tf.map_fn(
          fn=lambda x: self._resize_and_pad_masks(x[0], x[1]),
          elems=(
              batched_segmentation_masks,
              image_info),
          fn_output_signature=tf.float32,
          parallel_iterations=32)
    else:
      batched_segmentation_masks = tf.image.resize(
          batched_segmentation_masks,
          size=self._output_size,
          method='bilinear')

    batched_segmentation_masks = tf.expand_dims(tf.to_int32(
        tf.argmax(batched_segmentation_masks, axis=-1)), axis=-1)

    panoptic_masks = tf.map_fn(
        fn=lambda x: self._generate_panoptic_masks(  # pylint:disable=g-long-lambda
            x[0], x[1], x[2], x[3], x[4]),
        elems=(
            batched_boxes,
            batched_scores,
            batched_classes,
            batched_detections_masks,
            batched_segmentation_masks),
        fn_output_signature={
            'category_mask': tf.float32,
            'instance_mask': tf.float32
        }, parallel_iterations=32)

    for k, v in panoptic_masks.items():
      panoptic_masks[k] = tf.to_int32(v)

    return panoptic_masks
//  [python]
def call(self,
           inputs: tf.Tensor,
           image_info: Optional[tf.Tensor] = None) -> Dict[str, tf.Tensor]:
    detections = inputs

    batched_scores = detections['detection_scores']
    batched_classes = detections['detection_classes']
    batched_detections_masks = tf.expand_dims(
        detections['detection_masks'], axis=-1)
    batched_boxes = detections['detection_boxes']
    batched_segmentation_masks = tf.cast(
        detections['segmentation_outputs'], dtype=tf.float32)

    if self._rescale_predictions:
      scale = tf.tile(
          tf.cast(image_info[:, 2:3, :], dtype=batched_boxes.dtype),
          multiples=[1, 1, 2])
      batched_boxes /= scale

      batched_segmentation_masks = tf.map_fn(
          fn=lambda x: self._resize_and_pad_masks(x[0], x[1]),
          elems=(
              batched_segmentation_masks,
              image_info),
          fn_output_signature=tf.float32,
          parallel_iterations=32)
    else:
      batched_segmentation_masks = tf.image.resize(
          batched_segmentation_masks,
          size=self._output_size,
          method='bilinear')

    batched_segmentation_masks = tf.expand_dims(tf.cast(
        tf.argmax(batched_segmentation_masks, axis=-1),
        dtype=tf.float32), axis=-1)

    panoptic_masks = tf.map_fn(
        fn=lambda x: self._generate_panoptic_masks(  # pylint:disable=g-long-lambda
            x[0], x[1], x[2], x[3], x[4]),
        elems=(
            batched_boxes,
            batched_scores,
            batched_classes,
            batched_detections_masks,
            batched_segmentation_masks),
        fn_output_signature={
            'category_mask': tf.float32,
            'instance_mask': tf.float32
        }, parallel_iterations=32)

    for k, v in panoptic_masks.items():
      panoptic_masks[k] = tf.cast(v, dtype=tf.int32)

    return panoptic_masks
//  [python]
def loss_with_spring(logits_left,logits_right,labels):
    margin = 1.0
    #tf.cast(labels,tf.float32)
    
    labels_t = tf.to_float(labels)
    
    left_feature=tf.nn.l2_normalize(logits_left, dim = 1)
    right_feature=tf.nn.l2_normalize(logits_right, dim = 1)
    
    labels_f = tf.subtract(1.0, labels_t, name="1-yi")          # labels_ = !labels;
    eucd2 = tf.pow(tf.subtract(left_feature, right_feature), 2)
    eucd2 = tf.reduce_sum(eucd2, 1)
    eucd = tf.sqrt(eucd2+1e-6, name="eucd")
    C = tf.constant(margin, name="C")
    # yi*||CNN(p1i)-CNN(p2i)||^2 + (1-yi)*max(0, C-||CNN(p1i)-CNN(p2i)||^2)
    pos = tf.multiply(labels_t, eucd2, name="yi_x_eucd2")
    # neg = tf.multiply(labels_f, tf.subtract(0.0,eucd2), name="yi_x_eucd2")
    # neg = tf.multiply(labels_f, tf.maximum(0.0, tf.subtract(C,eucd2)), name="Nyi_x_C-eucd_xx_2")
    neg = tf.multiply(labels_f, tf.pow(tf.maximum(tf.subtract(C, eucd), 0), 2), name="Nyi_x_C-eucd_xx_2")
    losses = tf.add(pos, neg, name="losses")
    loss = tf.reduce_mean(losses, name="loss")/4
    return loss,eucd
//  [python]
def loss_with_spring(logits_left,logits_right,labels):
    margin = 1.0
    #tf.cast(labels,tf.float32)
    
    labels_t = tf.cast(labels,tf.float32)
    
    left_feature=tf.nn.l2_normalize(logits_left, dim = 1)
    right_feature=tf.nn.l2_normalize(logits_right, dim = 1)
    
    labels_f = tf.subtract(1.0, labels_t, name="1-yi")          # labels_ = !labels;
    eucd2 = tf.pow(tf.subtract(left_feature, right_feature), 2)
    eucd2 = tf.reduce_sum(eucd2, 1)
    eucd = tf.sqrt(eucd2+1e-6, name="eucd")
    C = tf.constant(margin, name="C")
    # yi*||CNN(p1i)-CNN(p2i)||^2 + (1-yi)*max(0, C-||CNN(p1i)-CNN(p2i)||^2)
    pos = tf.multiply(labels_t, eucd2, name="yi_x_eucd2")
    # neg = tf.multiply(labels_f, tf.subtract(0.0,eucd2), name="yi_x_eucd2")
    # neg = tf.multiply(labels_f, tf.maximum(0.0, tf.subtract(C,eucd2)), name="Nyi_x_C-eucd_xx_2")
    neg = tf.multiply(labels_f, tf.pow(tf.maximum(tf.subtract(C, eucd), 0), 2), name="Nyi_x_C-eucd_xx_2")
    losses = tf.add(pos, neg, name="losses")
    loss = tf.reduce_mean(losses, name="loss")/4
    return loss,eucd
//  [python]
def test_dtypes_flexible(self):
        """Test that encoder supports variable dtypes and
        yields same result."""

        dt_supported = (tf.float16, tf.float32, tf.float64, tf.int8,
            tf.int32, tf.int64, tf.uint8, tf.uint16, tf.uint32)

        bs = 10
        k = 32
        n = 64

        source = BinarySource()
        frozen_pos, _ = generate_5g_ranking(k, n)
        enc_ref = PolarEncoder(frozen_pos, n, dtype=tf.float32)

        u = source([bs, k])
        c_ref = enc_ref(u)

        for dt in dt_supported:
            enc = PolarEncoder(frozen_pos, n, dtype=dt)
            u_dt = tf.cast(u, dt)
            c = enc(u_dt)

            c_32 = tf.cast(c, tf.float32)

            self.assertTrue(np.array_equal(c_ref.numpy(), c_32.numpy()))
//  [python]
def test_dtypes_flexible(self):
        """Test that encoder supports variable dtypes and
        yields same result for all dtypes."""

        dt_supported = (tf.float16, tf.float32, tf.float64, tf.int8,
            tf.int32, tf.int64, tf.uint8, tf.uint16, tf.uint32)

        bs = 10
        k = 32
        n = 64

        source = BinarySource()
        enc_ref = Polar5GEncoder(k, n, dtype=tf.float32)

        u = source([bs, k])
        c_ref = enc_ref(u)

        for dt in dt_supported:
            enc = Polar5GEncoder(k, n, dtype=dt)
            u_dt = tf.cast(u, dt)
            c = enc(u_dt)

            c_32 = tf.cast(c, tf.float32)

            self.assertTrue(np.array_equal(c_ref.numpy(), c_32.numpy()))