jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def affine_to_dense_shift(matrix, shape, shift_center=True, warp_right=None):
    """
    Convert N-dimensional (ND) matrix transforms to dense displacement fields.

    Algorithm:
        1. Build and (optionally) shift grid to center of image.
        2. Apply matrices to each index coordinate.
        3. Subtract grid.

    Parameters:
        matrix: Affine matrix of shape (..., M, N + 1), where M is N or N + 1. Can have any batch
            dimensions.
        shape: ND shape of the output space.
        shift_center: Shift grid to image center.
        warp_right: Right-compose the matrix transform with a displacement field of shape
            (..., *shape, N), with batch dimensions broadcastable to those of `matrix`.

    Returns:
        Dense shift (warp) of shape (..., *shape, N).

    Notes:
        There used to be an argument for choosing between matrix ('ij') and Cartesian ('xy')
        indexing. Due to inconsistencies in how some functions and layers handled xy-indexing, we
        removed it in favor of default ij-indexing to minimize the potential for confusion.

    """
    if isinstance(shape, (tf.compat.v1.Dimension, tf.TensorShape)):
        shape = shape.as_list()

    if not tf.is_tensor(matrix) or not matrix.dtype.is_floating:
        matrix = tf.cast(matrix, tf.float32)

    # check input shapes
    ndims = len(shape)
    if matrix.shape[-1] != (ndims + 1):
        matdim = matrix.shape[-1] - 1
        raise ValueError(f'Affine ({matdim}D) does not match target shape ({ndims}D).')
    validate_affine_shape(matrix.shape)

    # coordinate grid
    mesh = (tf.range(s, dtype=matrix.dtype) for s in shape)
    if shift_center:
        mesh = (m - 0.5 * (s - 1) for m, s in zip(mesh, shape))
    mesh = [tf.reshape(m, shape=(-1,)) for m in tf.meshgrid(*mesh, indexing='ij')]
    mesh = tf.stack(mesh)  # N x nb_voxels
    out = mesh

    # optionally right-compose with warp field
    if warp_right is not None:
        if not tf.is_tensor(warp_right) or warp_right.dtype != matrix.dtype:
            warp_right = tf.cast(warp_right, matrix.dtype)
        flat_shape = tf.concat((tf.shape(warp_right)[:-1 - ndims], (-1, ndims)), axis=0)
        warp_right = tf.reshape(warp_right, flat_shape)  # ... x nb_voxels x N
        out += tf.linalg.matrix_transpose(warp_right)  # ... x N x nb_voxels

    # compute locations, subtract grid to obtain shift
    out = matrix[..., :ndims, :-1] @ out + matrix[..., :ndims, -1:]  # ... x N x nb_voxels
    out = tf.linalg.matrix_transpose(out - mesh)  # ... x nb_voxels x N

    # restore shape
    shape = tf.concat((tf.shape(matrix)[:-2], (*shape, ndims)), axis=0)
    return tf.reshape(out, shape)
//  [python]
def __calc_prune_ratio_dyn(self, prune_ratio_fnl):
    """Calculate the dynamic pruning ratio.

    Args:
    * prune_ratio_fnl: final pruning ratio

    Returns:
    * prune_ratio_dyn: dynamic pruning ratio
    """

    idx_iter_beg = int(self.nb_iters_train * FLAGS.ws_iter_ratio_beg)
    idx_iter_end = int(self.nb_iters_train * FLAGS.ws_iter_ratio_end)
    base = t.cast(self.global_step - idx_iter_beg, t.float32) / (idx_iter_end - idx_iter_beg)
    base = tf.minimum(1.0, tf.maximum(0.0, base))
    prune_ratio_dyn = prune_ratio_fnl * (1.0 - tf.pow(1.0 - base, FLAGS.ws_prune_ratio_exp))

    return prune_ratio_dyn
//  [python]
def __calc_prune_ratio_dyn(self, prune_ratio_fnl):
    """Calculate the dynamic pruning ratio.

    Args:
    * prune_ratio_fnl: final pruning ratio

    Returns:
    * prune_ratio_dyn: dynamic pruning ratio
    """

    idx_iter_beg = int(self.nb_iters_train * FLAGS.ws_iter_ratio_beg)
    idx_iter_end = int(self.nb_iters_train * FLAGS.ws_iter_ratio_end)
    base = tf.cast(self.global_step - idx_iter_beg, tf.float32) / (idx_iter_end - idx_iter_beg)
    base = tf.minimum(1.0, tf.maximum(0.0, base))
    prune_ratio_dyn = prune_ratio_fnl * (1.0 - tf.pow(1.0 - base, FLAGS.ws_prune_ratio_exp))

    return prune_ratio_dyn
//  [python]
def call(self, g, x, etypes, norm=None):
        """Forward computation

        Parameters
        ----------
        g : DGLGraph
            The graph.
        x : tf.Tensor
            Input node features. Could be either

                * :math:`(|V|, D)` dense tensor
                * :math:`(|V|,)` int64 vector, representing the categorical values of each
                  node. We then treat the input feature as an one-hot encoding feature.
        etypes : tf.Tensor
            Edge type tensor. Shape: :math:`(|E|,)`
        norm : tf.Tensor
            Optional edge normalizer tensor. Shape: :math:`(|E|, 1)`

        Returns
        -------
        tf.Tensor
            New node features.
        """
        assert g.is_homogeneous, (
            "not a homogeneous graph; convert it with to_homogeneous "
            "and pass in the edge type as argument"
        )
        with g.local_scope():
            g.ndata["h"] = x
            g.edata["type"] = t.cast(etypes, t.int64)
            if norm is not None:
                g.edata["norm"] = norm
            if self.self_loop:
                loop_message = utils.matmul_maybe_select(x, self.loop_weight)
            # message passing
            g.update_all(self.message_func, fn.sum(msg="msg", out="h"))
            # apply bias and activation
            node_repr = g.ndata["h"]
            if self.bias:
                node_repr = node_repr + self.h_bias
            if self.self_loop:
                node_repr = node_repr + loop_message
            if self.activation:
                node_repr = self.activation(node_repr)
            node_repr = self.dropout(node_repr)
            return node_repr
//  [python]
def call(self, g, x, etypes, norm=None):
        """Forward computation

        Parameters
        ----------
        g : DGLGraph
            The graph.
        x : tf.Tensor
            Input node features. Could be either

                * :math:`(|V|, D)` dense tensor
                * :math:`(|V|,)` int64 vector, representing the categorical values of each
                  node. We then treat the input feature as an one-hot encoding feature.
        etypes : tf.Tensor
            Edge type tensor. Shape: :math:`(|E|,)`
        norm : tf.Tensor
            Optional edge normalizer tensor. Shape: :math:`(|E|, 1)`

        Returns
        -------
        tf.Tensor
            New node features.
        """
        assert g.is_homogeneous, (
            "not a homogeneous graph; convert it with to_homogeneous "
            "and pass in the edge type as argument"
        )
        with g.local_scope():
            g.ndata["h"] = x
            g.edata["type"] = tf.cast(etypes, tf.int64)
            if norm is not None:
                g.edata["norm"] = norm
            if self.self_loop:
                loop_message = utils.matmul_maybe_select(x, self.loop_weight)
            # message passing
            g.update_all(self.message_func, fn.sum(msg="msg", out="h"))
            # apply bias and activation
            node_repr = g.ndata["h"]
            if self.bias:
                node_repr = node_repr + self.h_bias
            if self.self_loop:
                node_repr = node_repr + loop_message
            if self.activation:
                node_repr = self.activation(node_repr)
            node_repr = self.dropout(node_repr)
            return node_repr
//  [python]
def get_label_weight_mask(labels, ignore_label, num_classes, label_weights=1.0):
  """Gets the label weight mask.

  Args:
    labels: A Tensor of labels with the shape of [-1].
    ignore_label: Integer, label to ignore.
    num_classes: Integer, the number of semantic classes.
    label_weights: A float or a list of weights. If it is a float, it means all
      the labels have the same weight. If it is a list of weights, then each
      element in the list represents the weight for the label of its index, for
      example, label_weights = [0.1, 0.5] means the weight for label 0 is 0.1
      and the weight for label 1 is 0.5.

  Returns:
    A Tensor of label weights with the same shape of labels, each element is the
      weight for the label with the same index in labels and the element is 0.0
      if the label is to ignore.

  Raises:
    ValueError: If label_weights is neither a float nor a list, or if
      label_weights is a list and its length is not equal to num_classes.
  """
  if not isinstance(label_weights, (float, list)):
    raise ValueError(
        'The type of label_weights is invalid, it must be a float or a list.')

  if isinstance(label_weights, list) and len(label_weights) != num_classes:
    raise ValueError(
        'Length of label_weights must be equal to num_classes if it is a list, '
        'label_weights: %s, num_classes: %d.' % (label_weights, num_classes))

  not_ignore_mask = tf.not_equal(labels, ignore_label)
  not_ignore_mask = tf.to_float(not_ignore_mask)
  if isinstance(label_weights, float):
    return not_ignore_mask * label_weights

  label_weights = tf.constant(label_weights, tf.float32)
  weight_mask = tf.einsum('...y,y->...',
                          tf.one_hot(labels, num_classes, dtype=tf.float32),
                          label_weights)
  return tf.multiply(not_ignore_mask, weight_mask)
//  [python]
def get_label_weight_mask(labels, ignore_label, num_classes, label_weights=1.0):
  """Gets the label weight mask.

  Args:
    labels: A Tensor of labels with the shape of [-1].
    ignore_label: Integer, label to ignore.
    num_classes: Integer, the number of semantic classes.
    label_weights: A float or a list of weights. If it is a float, it means all
      the labels have the same weight. If it is a list of weights, then each
      element in the list represents the weight for the label of its index, for
      example, label_weights = [0.1, 0.5] means the weight for label 0 is 0.1
      and the weight for label 1 is 0.5.

  Returns:
    A Tensor of label weights with the same shape of labels, each element is the
      weight for the label with the same index in labels and the element is 0.0
      if the label is to ignore.

  Raises:
    ValueError: If label_weights is neither a float nor a list, or if
      label_weights is a list and its length is not equal to num_classes.
  """
  if not isinstance(label_weights, (float, list)):
    raise ValueError(
        'The type of label_weights is invalid, it must be a float or a list.')

  if isinstance(label_weights, list) and len(label_weights) != num_classes:
    raise ValueError(
        'Length of label_weights must be equal to num_classes if it is a list, '
        'label_weights: %s, num_classes: %d.' % (label_weights, num_classes))

  not_ignore_mask = tf.not_equal(labels, ignore_label)
  not_ignore_mask = tf.cast(not_ignore_mask, tf.float32)
  if isinstance(label_weights, float):
    return not_ignore_mask * label_weights

  label_weights = tf.constant(label_weights, tf.float32)
  weight_mask = tf.einsum('...y,y->...',
                          tf.one_hot(labels, num_classes, dtype=tf.float32),
                          label_weights)
  return tf.multiply(not_ignore_mask, weight_mask)
//  [python]
def __build_masks(self):
    """build masks and corresponding operations for weight sparsification.

    Returns:
    * masks: list of masks for weight sparsification
    * prune_op: pruning operation
    """

    masks, prune_ops = [], []
    with tf.variable_scope(self.mask_scope):
      for var, var_name_n_prune_ratio in zip(self.maskable_vars, self.var_names_n_prune_ratios):
        # obtain the dynamic pruning ratio
        assert var.name == var_name_n_prune_ratio[0], \
            'unmatched variable names: %s vs. %s' % (var.name, var_name_n_prune_ratio[0])
        prune_ratio = self.__calc_prune_ratio_dyn(var_name_n_prune_ratio[1])

        # create a mask and non-masked backup for each variable
        name = var.name.replace(':0', '_mask')
        mask = tf.get_variable(name, initializer=tf.ones(var.shape), trainable=False)
        name = var.name.replace(':0', '_var_bkup')
        var_bkup = tf.get_variable(name, initializer=var.initialized_value(), trainable=False)

        # create update operations
        var_bkup_update_op = var_bkup.assign(tf.where(mask > 0.5, var, var_bkup))
        with tf.control_dependencies([var_bkup_update_op]):
          mask_thres = tf.contrib.distributions.percentile(tf.abs(var_bkup), prune_ratio * 100)
          mask_update_op = mask.assign(tf.to_float(tf.abs(var_bkup) > mask_thres))
        with tf.control_dependencies([mask_update_op]):
          prune_op = var.assign(var_bkup * mask)

        # record pruning masks & operations
        masks += [mask]
        prune_ops += [prune_op]

    return masks, tf.group(prune_ops)
//  [python]
def __build_masks(self):
    """build masks and corresponding operations for weight sparsification.

    Returns:
    * masks: list of masks for weight sparsification
    * prune_op: pruning operation
    """

    masks, prune_ops = [], []
    with tf.variable_scope(self.mask_scope):
      for var, var_name_n_prune_ratio in zip(self.maskable_vars, self.var_names_n_prune_ratios):
        # obtain the dynamic pruning ratio
        assert var.name == var_name_n_prune_ratio[0], \
            'unmatched variable names: %s vs. %s' % (var.name, var_name_n_prune_ratio[0])
        prune_ratio = self.__calc_prune_ratio_dyn(var_name_n_prune_ratio[1])

        # create a mask and non-masked backup for each variable
        name = var.name.replace(':0', '_mask')
        mask = tf.get_variable(name, initializer=tf.ones(var.shape), trainable=False)
        name = var.name.replace(':0', '_var_bkup')
        var_bkup = tf.get_variable(name, initializer=var.initialized_value(), trainable=False)

        # create update operations
        var_bkup_update_op = var_bkup.assign(tf.where(mask > 0.5, var, var_bkup))
        with tf.control_dependencies([var_bkup_update_op]):
          mask_thres = tf.contrib.distributions.percentile(tf.abs(var_bkup), prune_ratio * 100)
          mask_update_op = mask.assign(tf.cast(tf.abs(var_bkup) > mask_thres, tf.float32))
        with tf.control_dependencies([mask_update_op]):
          prune_op = var.assign(var_bkup * mask)

        # record pruning masks & operations
        masks += [mask]
        prune_ops += [prune_op]

    return masks, tf.group(prune_ops)
//  [python]
def trim_zeros_graph(boxes, name=None):
    """Often boxes are represented with matrices of shape [N, 4] and
    are padded with zeros. This removes zero boxes.

    boxes: [N, 4] matrix of boxes.
    non_zeros: [N] a 1D boolean mask identifying the rows to keep
    """
    non_zeros = tf.cast(tf.reduce_sum(tf.abs(boxes), axis=1), tf.bool)
    boxes = tf.boolean_mask(boxes, non_zeros, name=name)
    return boxes, non_zeros
//  [python]
def Loss_SpatialColorSmooth_add_SelfContain(Z,X,gamma=1e-1,knn=10):
    '''
    function to return spatial and color smoothness constraint loss
    :param Z: Input point cloud feature embedding float B*N*D
    :param X: Input point cloud  float B*N*6    XYZRGB
    :param W: Input knn weight matrix  float B*N*Knn
    :param Ind: Input knn index  int B*N*Knn
    :return: spatial smooth loss
    '''

    ## Compute Full Weight Matrix
    N = tf.shape(X)[1]

    # XYZ channel
    X_2 = tf.tile(tf.reduce_sum(X ** 2, axis=-1, keepdims=True), multiples=[1, 1, N])  # B*N*N
    Y_2 = tf.einsum('ijk->ikj', X_2)
    XY = tf.einsum('ijk,ilk->ijl', X, X)
    Dmat = X_2 + Y_2 - 2 * XY  # B*N*N
    Dmat = tf.to_float(tf.cast(Dmat < 0, tf.int32)) * tf.zeros_like(Dmat) + tf.to_float(tf.cast(Dmat > 0, tf.int32)) * Dmat

    # W = tf.exp(-Dmat / gamma)    # B*N*N

    ## Compute Knn Graph
    # Ind = tf.argsort(Dmat)[:,:,0:knn]   # B*N*knn
    _, Ind = tf.nn.top_k(-Dmat, k = knn, sorted=True) # B*N*knn
    W_knn = []
    for d_i, Ind_i in zip(tf.unstack(Dmat,axis=0), tf.unstack(Ind)):
        W_knn.append(tf.batch_gather(tf.exp(-d_i / gamma), Ind_i))
    W_knn = tf.stack(W_knn)     # B*N*knn

    #### Compute Loss
    Z_tilde = Tool.batch_gather_v1(Z,Ind)   # indexed X B*N*Knn*D
    Z_expand = tf.tile(Z[:,:,tf.newaxis,:],[1,1,knn,1])    # expanded X B*N*Knn*D
    loss = W_knn*tf.reduce_mean((Z_expand - Z_tilde)**2,axis=-1) # B*N*Knn
    # loss = tf.reduce_sum(loss)/tf.cast(tf.shape(loss)[0]*tf.shape(loss)[1]*tf.shape(loss)[2],tf.float32)
    loss = tf.reduce_mean(loss)

    return loss
//  [python]
def Loss_SpatialColorSmooth_add_SelfContain(Z,X,gamma=1e-1,knn=10):
    '''
    function to return spatial and color smoothness constraint loss
    :param Z: Input point cloud feature embedding float B*N*D
    :param X: Input point cloud  float B*N*6    XYZRGB
    :param W: Input knn weight matrix  float B*N*Knn
    :param Ind: Input knn index  int B*N*Knn
    :return: spatial smooth loss
    '''

    ## Compute Full Weight Matrix
    N = tf.shape(X)[1]

    # XYZ channel
    X_2 = tf.tile(tf.reduce_sum(X ** 2, axis=-1, keepdims=True), multiples=[1, 1, N])  # B*N*N
    Y_2 = tf.einsum('ijk->ikj', X_2)
    XY = tf.einsum('ijk,ilk->ijl', X, X)
    Dmat = X_2 + Y_2 - 2 * XY  # B*N*N
    Dmat = tf.cast(Dmat < 0, tf.float32) * tf.zeros_like(Dmat) + tf.cast(Dmat > 0, tf.float32) * Dmat

    # W = tf.exp(-Dmat / gamma)    # B*N*N

    ## Compute Knn Graph
    # Ind = tf.argsort(Dmat)[:,:,0:knn]   # B*N*knn
    _, Ind = tf.nn.top_k(-Dmat, k = knn, sorted=True) # B*N*knn
    W_knn = []
    for d_i, Ind_i in zip(tf.unstack(Dmat,axis=0), tf.unstack(Ind)):
        W_knn.append(tf.batch_gather(tf.exp(-d_i / gamma), Ind_i))
    W_knn = tf.stack(W_knn)     # B*N*knn

    #### Compute Loss
    Z_tilde = Tool.batch_gather_v1(Z,Ind)   # indexed X B*N*Knn*D
    Z_expand = tf.tile(Z[:,:,tf.newaxis,:],[1,1,knn,1])    # expanded X B*N*Knn*D
    loss = W_knn*tf.reduce_mean((Z_expand - Z_tilde)**2,axis=-1) # B*N*Knn
    # loss = tf.reduce_sum(loss)/tf.cast(tf.shape(loss)[0]*tf.shape(loss)[1]*tf.shape(loss)[2],tf.float32)
    loss = tf.reduce_mean(loss)

    return loss
//  [python]
def read_data(file_q):
    # Code from https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_input.py
    class CIFAR10Record(object):
        pass

    result = CIFAR10Record()

    # Dimensions of the images in the CIFAR-10 dataset.
    # See http://www.cs.toronto.edu/~kriz/cifar.html for a description of the
    # input format.
    label_bytes = 1  # 2 for CIFAR-100
    result.height = 32
    result.width = 32
    result.depth = 3
    image_bytes = result.height * result.width * result.depth
    # Every record consists of a label followed by the image, with a
    # fixed number of bytes for each.
    record_bytes = label_bytes + image_bytes

    # Read a record, getting filenames from the filename_queue.  No
    # header or footer in the CIFAR-10 format, so we leave header_bytes
    # and footer_bytes at their default of 0.
    reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
    result.key, value = reader.read(file_q)

    # Convert from a string to a vector of uint8 that is record_bytes long.
    record_bytes = tf.decode_raw(value, tf.uint8)

    # The first bytes represent the label, which we convert from uint8->int32.
    result.label = tf.cast(
        tf.strided_slice(record_bytes, [0], [label_bytes]), tf.int32)

    # The remaining bytes after the label represent the image, which we reshape
    # from [depth * height * width] to [depth, height, width].
    depth_major = tf.reshape(
        tf.strided_slice(record_bytes, [label_bytes],
                         [label_bytes + image_bytes]),
        [result.depth, result.height, result.width])
    # Convert from [depth, height, width] to [height, width, depth].
    result.uint8image = tf.transpose(depth_major, [1, 2, 0])

    reshaped_image = tf.to_float(result.uint8image)

    height = 24
    width = 24

    # Image processing for evaluation.
    # Crop the central [height, width] of the image.
    resized_image = tf.image.resize_image_with_crop_or_pad(reshaped_image,
                                                           height, width)

    # Subtract off the mean and divide by the variance of the pixels.
    float_image = tf.image.per_image_standardization(resized_image)

    # Set the shapes of tensors.
    float_image.set_shape([height, width, 3])
    result.label.set_shape([1])

    return float_image, result.label
//  [python]
def read_data(file_q):
    # Code from https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_input.py
    class CIFAR10Record(object):
        pass

    result = CIFAR10Record()

    # Dimensions of the images in the CIFAR-10 dataset.
    # See http://www.cs.toronto.edu/~kriz/cifar.html for a description of the
    # input format.
    label_bytes = 1  # 2 for CIFAR-100
    result.height = 32
    result.width = 32
    result.depth = 3
    image_bytes = result.height * result.width * result.depth
    # Every record consists of a label followed by the image, with a
    # fixed number of bytes for each.
    record_bytes = label_bytes + image_bytes

    # Read a record, getting filenames from the filename_queue.  No
    # header or footer in the CIFAR-10 format, so we leave header_bytes
    # and footer_bytes at their default of 0.
    reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
    result.key, value = reader.read(file_q)

    # Convert from a string to a vector of uint8 that is record_bytes long.
    record_bytes = tf.decode_raw(value, tf.uint8)

    # The first bytes represent the label, which we convert from uint8->int32.
    result.label = tf.cast(
        tf.strided_slice(record_bytes, [0], [label_bytes]), tf.int32)

    # The remaining bytes after the label represent the image, which we reshape
    # from [depth * height * width] to [depth, height, width].
    depth_major = tf.reshape(
        tf.strided_slice(record_bytes, [label_bytes],
                         [label_bytes + image_bytes]),
        [result.depth, result.height, result.width])
    # Convert from [depth, height, width] to [height, width, depth].
    result.uint8image = tf.transpose(depth_major, [1, 2, 0])

    reshaped_image = tf.cast(result.uint8image, tf.float32)

    height = 24
    width = 24

    # Image processing for evaluation.
    # Crop the central [height, width] of the image.
    resized_image = tf.image.resize_image_with_crop_or_pad(reshaped_image,
                                                           height, width)

    # Subtract off the mean and divide by the variance of the pixels.
    float_image = tf.image.per_image_standardization(resized_image)

    # Set the shapes of tensors.
    float_image.set_shape([height, width, 3])
    result.label.set_shape([1])

    return float_image, result.label
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

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

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

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

    # Smooth-L1 Loss
    loss = K.switch(tf.size(target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

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

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

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

    # Smooth-L1 Loss
    loss = K.switch(tf.size(target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def Loss_SpatialColorSmoothAdd_UnknownBatch_SelfContain(Z,X,gamma=1e-1,knn=10):
    '''
    function to return spatial and color smoothness constraint loss
    :param Z: Input point cloud feature embedding float B*N*D
    :param X: Input point cloud  float B*N*6    XYZRGB
    :param W: Input knn weight matrix  float B*N*Knn
    :param Ind: Input knn index  int B*N*Knn
    :return: spatial smooth loss
    '''

    ## Compute Full Weight Matrix
    N = tf.shape(X)[1]

    # XYZ channel
    X_2 = tf.tile(tf.reduce_sum(X ** 2, axis=-1, keepdims=True), multiples=[1, 1, N])  # B*N*N
    Y_2 = tf.einsum('ijk->ikj', X_2)
    XY = tf.einsum('ijk,ilk->ijl', X, X)
    Dmat = X_2 + Y_2 - 2 * XY  # B*N*N
    Dmat = tf.to_float(tf.cast(Dmat < 0, tf.int32)) * tf.zeros_like(Dmat) + tf.to_float(tf.cast(Dmat > 0, tf.int32)) * Dmat

    # W = tf.exp(-Dmat / gamma)    # B*N*N
    def custom_gather_v2(a, b):
        def apply_gather(x):
            return tf.gather(x[0], tf.cast(x[1], tf.int32))

        a = tf.cast(a, dtype=tf.float32)
        b = tf.cast(b, dtype=tf.float32)
        stacked = tf.stack([a, b], axis=1)
        gathered = tf.map_fn(apply_gather, stacked)
        return tf.stack(gathered, axis=0)
    
    ## Compute Knn Graph
    # Ind = tf.argsort(Dmat)[:,:,0:knn]   # B*N*knn
    negDmat_knn, Ind = tf.nn.top_k(-Dmat, k = knn, sorted=True) # B*N*knn
    W_knn = tf.exp(negDmat_knn/gamma)
    # W_knn = []
    #
    # W_knn = custom_gather_v2(tf.exp(-Dmat/gamma), Ind)
    
    # for d_i, Ind_i in zip(tf.unstack(Dmat,axis=0), tf.unstack(Ind)):
    #     W_knn.append(tf.batch_gather(tf.exp(-d_i / gamma), Ind_i))
    # W_knn = tf.stack(W_knn)     # B*N*knn

    #### Compute Loss
    Z_tilde = Tool.batch_gather_v1(Z,Ind)   # indexed X B*N*Knn*D
    Z_expand = tf.tile(Z[:,:,tf.newaxis,:],[1,1,knn,1])    # expanded X B*N*Knn*D
    loss = W_knn*tf.reduce_sum((Z_expand - Z_tilde)**2,axis=-1) # B*N*Knn
    # loss = tf.reduce_sum(loss)/tf.cast(tf.shape(loss)[0]*tf.shape(loss)[1]*tf.shape(loss)[2],tf.float32)
    loss = tf.reduce_mean(loss)

    return loss
//  [python]
def Loss_SpatialColorSmoothAdd_UnknownBatch_SelfContain(Z,X,gamma=1e-1,knn=10):
    '''
    function to return spatial and color smoothness constraint loss
    :param Z: Input point cloud feature embedding float B*N*D
    :param X: Input point cloud  float B*N*6    XYZRGB
    :param W: Input knn weight matrix  float B*N*Knn
    :param Ind: Input knn index  int B*N*Knn
    :return: spatial smooth loss
    '''

    ## Compute Full Weight Matrix
    N = tf.shape(X)[1]

    # XYZ channel
    X_2 = tf.tile(tf.reduce_sum(X ** 2, axis=-1, keepdims=True), multiples=[1, 1, N])  # B*N*N
    Y_2 = tf.einsum('ijk->ikj', X_2)
    XY = tf.einsum('ijk,ilk->ijl', X, X)
    Dmat = X_2 + Y_2 - 2 * XY  # B*N*N
    Dmat = tf.cast(Dmat < 0, tf.float32) * tf.zeros_like(Dmat) + tf.cast(Dmat > 0, tf.float32) * Dmat

    # W = tf.exp(-Dmat / gamma)    # B*N*N
    def custom_gather_v2(a, b):
        def apply_gather(x):
            return tf.gather(x[0], tf.cast(x[1], tf.int32))

        a = tf.cast(a, dtype=tf.float32)
        b = tf.cast(b, dtype=tf.float32)
        stacked = tf.stack([a, b], axis=1)
        gathered = tf.map_fn(apply_gather, stacked)
        return tf.stack(gathered, axis=0)
    
    ## Compute Knn Graph
    # Ind = tf.argsort(Dmat)[:,:,0:knn]   # B*N*knn
    negDmat_knn, Ind = tf.nn.top_k(-Dmat, k = knn, sorted=True) # B*N*knn
    W_knn = tf.exp(negDmat_knn/gamma)
    # W_knn = []
    #
    # W_knn = custom_gather_v2(tf.exp(-Dmat/gamma), Ind)
    
    # for d_i, Ind_i in zip(tf.unstack(Dmat,axis=0), tf.unstack(Ind)):
    #     W_knn.append(tf.batch_gather(tf.exp(-d_i / gamma), Ind_i))
    # W_knn = tf.stack(W_knn)     # B*N*knn

    #### Compute Loss
    Z_tilde = Tool.batch_gather_v1(Z,Ind)   # indexed X B*N*Knn*D
    Z_expand = tf.tile(Z[:,:,tf.newaxis,:],[1,1,knn,1])    # expanded X B*N*Knn*D
    loss = W_knn*tf.reduce_sum((Z_expand - Z_tilde)**2,axis=-1) # B*N*Knn
    # loss = tf.reduce_sum(loss)/tf.cast(tf.shape(loss)[0]*tf.shape(loss)[1]*tf.shape(loss)[2],tf.float32)
    loss = tf.reduce_mean(loss)

    return loss
//  [python]
def generate_queries_random(self, number_queries):
    """Generate rotation matrices from SO(3) uniformly at random.

    Args:
      number_queries: The number of queries.
    Returns:
      A tensor of shape [number_queries, 3, 3].
    """
    random_quaternions = tfg.quaternion.normalized_random_uniform(
        (number_queries,))
    random_rotations = tfg.rotation_matrix_3d.from_quaternion(
        random_quaternions)
    random_rotations = tf.cast(random_rotations, tf.float32)
    return random_rotations
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        
        optimizer = keras.optimizers.Adam(
            lr=learning_rate,
            clipnorm=self.config.GRADIENT_CLIP_NORM)

        """
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=self.config.GRADIENT_CLIP_NORM)
        """
        
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

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

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

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def _int_dataset_range(n):
  return tf.data.Dataset.range(n).map(lambda x: tf.to_int32(x))
//  [python]
def _int_dataset_range(n):
  return tf.data.Dataset.range(n).map(lambda x: tf.cast(x, tf.int32))
//  [python]
def preprocess_img(img_plac, gtbox_plac):
    '''

    :param img_plac: [H, W, 3] uint 8 img. In RGB.
    :param gtbox_plac: shape of [-1, 5]. [xmin, ymin, xmax, ymax, label]
    :return:
    '''

    img = tf.cast(img_plac, tf.float32)

    # gtboxes_and_label = tf.cast(gtbox_plac, tf.float32)
    img, gtboxes_and_label = image_preprocess.short_side_resize(img_tensor=img,
                                                                gtboxes_and_label=gtbox_plac,
                                                                target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
                                                                length_limitation=cfgs.IMG_MAX_LENGTH)
    img, gtboxes_and_label = image_preprocess.random_flip_left_right(img_tensor=img,
                                                                     gtboxes_and_label=gtboxes_and_label)
    img = img - tf.constant([[cfgs.PIXEL_MEAN]])
    img_batch = tf.expand_dims(img, axis=0)

    # gtboxes_and_label = tf.Print(gtboxes_and_label, [tf.shape(gtboxes_and_label)], message='gtbox shape')
    return img_batch, gtboxes_and_label
//  [python]
def create_normalizer_update(self, vector_input):
        mean_current_observation = tf.reduce_mean(vector_input, axis=0)
        new_mean = self.running_mean + (mean_current_observation - self.running_mean) / \
                   tf.to_float(tf.add(self.global_step, 1))
        new_variance = self.running_variance + (mean_current_observation - new_mean) * \
                       (mean_current_observation - self.running_mean)
        update_mean = tf.assign(self.running_mean, new_mean)
        update_variance = tf.assign(self.running_variance, new_variance)
        return update_mean, update_variance
//  [python]
def create_normalizer_update(self, vector_input):
        mean_current_observation = tf.reduce_mean(vector_input, axis=0)
        new_mean = self.running_mean + (mean_current_observation - self.running_mean) / \
                   tf.cast(tf.add(self.global_step, 1), tf.float32)
        new_variance = self.running_variance + (mean_current_observation - new_mean) * \
                       (mean_current_observation - self.running_mean)
        update_mean = tf.assign(self.running_mean, new_mean)
        update_variance = tf.assign(self.running_variance, new_variance)
        return update_mean, update_variance
//  [python]
def read_and_decode(filename_queue):
  reader = tf.TFRecordReader()
  _, serialized_example = reader.read(filename_queue)
  features = tf.parse_single_example(
      serialized_example,
      # Defaults are not specified since both keys are required.
      features={
          'image_raw': tf.FixedLenFeature([], tf.string),
          'label': tf.FixedLenFeature([], tf.int64),
      })

  # Convert from a scalar string tensor (whose single string has
  # length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
  # [mnist.IMAGE_PIXELS].
  image = tf.decode_raw(features['image_raw'], tf.uint8)
  image.set_shape([mnist.IMAGE_PIXELS])

  # OPTIONAL: Could reshape into a 28x28 image and apply distortions
  # here.  Since we are not applying any distortions in this
  # example, and the next step expects the image to be flattened
  # into a vector, we don't bother.

  # Convert from [0, 255] -> [-0.5, 0.5] floats.
  image = tf.to_float(image) * (1. / 255) - 0.5

  # Convert label from a scalar uint8 tensor to an int32 scalar.
  label = tf.to_int32(features['label'])

  return image, label
//  [python]
def read_and_decode(filename_queue):
  reader = tf.TFRecordReader()
  _, serialized_example = reader.read(filename_queue)
  features = tf.parse_single_example(
      serialized_example,
      # Defaults are not specified since both keys are required.
      features={
          'image_raw': tf.FixedLenFeature([], tf.string),
          'label': tf.FixedLenFeature([], tf.int64),
      })

  # Convert from a scalar string tensor (whose single string has
  # length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
  # [mnist.IMAGE_PIXELS].
  image = tf.decode_raw(features['image_raw'], tf.uint8)
  image.set_shape([mnist.IMAGE_PIXELS])

  # OPTIONAL: Could reshape into a 28x28 image and apply distortions
  # here.  Since we are not applying any distortions in this
  # example, and the next step expects the image to be flattened
  # into a vector, we don't bother.

  # Convert from [0, 255] -> [-0.5, 0.5] floats.
  image = tf.cast(image, tf.float32) * (1. / 255) - 0.5

  # Convert label from a scalar uint8 tensor to an int32 scalar.
  label = tf.cast(features['label'], tf.int32)

  return image, label
//  [python]
def create_vector_input(self, name='vector_observation'):
        """
        Creates ops for vector observation input.
        :param name: Name of the placeholder op.
        :param vec_obs_size: Size of stacked vector observation.
        :return:
        """
        self.vector_in = tf.placeholder(shape=[None, self.vec_obs_size], dtype=tf.float32,
                                        name=name)
        if self.normalize:
            self.running_mean = tf.get_variable("running_mean", [self.vec_obs_size],
                                                trainable=False, dtype=tf.float32,
                                                initializer=tf.zeros_initializer())
            self.running_variance = tf.get_variable("running_variance", [self.vec_obs_size],
                                                    trainable=False,
                                                    dtype=tf.float32,
                                                    initializer=tf.ones_initializer())
            self.update_mean, self.update_variance = self.create_normalizer_update(self.vector_in)

            self.normalized_state = tf.clip_by_value((self.vector_in - self.running_mean) / tf.sqrt(
                self.running_variance / (tf.to_float(self.global_step) + 1)), -5, 5,
                                                     name="normalized_state")
            return self.normalized_state
        else:
            return self.vector_in
//  [python]
def create_vector_input(self, name='vector_observation'):
        """
        Creates ops for vector observation input.
        :param name: Name of the placeholder op.
        :param vec_obs_size: Size of stacked vector observation.
        :return:
        """
        self.vector_in = tf.placeholder(shape=[None, self.vec_obs_size], dtype=tf.float32,
                                        name=name)
        if self.normalize:
            self.running_mean = tf.get_variable("running_mean", [self.vec_obs_size],
                                                trainable=False, dtype=tf.float32,
                                                initializer=tf.zeros_initializer())
            self.running_variance = tf.get_variable("running_variance", [self.vec_obs_size],
                                                    trainable=False,
                                                    dtype=tf.float32,
                                                    initializer=tf.ones_initializer())
            self.update_mean, self.update_variance = self.create_normalizer_update(self.vector_in)

            self.normalized_state = tf.clip_by_value((self.vector_in - self.running_mean) / tf.sqrt(
                self.running_variance / (tf.cast(self.global_step, tf.float32) + 1)), -5, 5,
                                                     name="normalized_state")
            return self.normalized_state
        else:
            return self.vector_in
//  [python]
def __init__(self, m_size, normalize, use_recurrent, brain, seed):
        tf.set_random_seed(seed)
        self.brain = brain
        self.vector_in = None
        self.global_step, self.increment_step = self.create_global_steps()
        self.visual_in = []
        self.batch_size = tf.placeholder(shape=None, dtype=tf.int32, name='batch_size')
        self.sequence_length = tf.placeholder(shape=None, dtype=tf.int32, name='sequence_length')
        self.mask_input = tf.placeholder(shape=[None], dtype=tf.float32, name='masks')
        self.mask = tf.to_int32(self.mask_input)
        self.use_recurrent = use_recurrent
        if self.use_recurrent:
            self.m_size = m_size
        else:
            self.m_size = 0
        self.normalize = normalize
        self.act_size = brain.vector_action_space_size
        self.vec_obs_size = brain.vector_observation_space_size * \
                            brain.num_stacked_vector_observations
        self.vis_obs_size = brain.number_visual_observations
        tf.Variable(int(brain.vector_action_space_type == 'continuous'),
                    name='is_continuous_control', trainable=False, dtype=tf.int32)
        tf.Variable(self._version_number_, name='version_number', trainable=False, dtype=tf.int32)
        tf.Variable(self.m_size, name="memory_size", trainable=False, dtype=tf.int32)
        if brain.vector_action_space_type == 'continuous':
            tf.Variable(self.act_size[0], name="action_output_shape", trainable=False, dtype=tf.int32)
        else:
            tf.Variable(sum(self.act_size), name="action_output_shape", trainable=False, dtype=tf.int32)
//  [python]
def __init__(self, m_size, normalize, use_recurrent, brain, seed):
        tf.set_random_seed(seed)
        self.brain = brain
        self.vector_in = None
        self.global_step, self.increment_step = self.create_global_steps()
        self.visual_in = []
        self.batch_size = tf.placeholder(shape=None, dtype=tf.int32, name='batch_size')
        self.sequence_length = tf.placeholder(shape=None, dtype=tf.int32, name='sequence_length')
        self.mask_input = tf.placeholder(shape=[None], dtype=tf.float32, name='masks')
        self.mask = tf.cast(self.mask_input, tf.int32)
        self.use_recurrent = use_recurrent
        if self.use_recurrent:
            self.m_size = m_size
        else:
            self.m_size = 0
        self.normalize = normalize
        self.act_size = brain.vector_action_space_size
        self.vec_obs_size = brain.vector_observation_space_size * \
                            brain.num_stacked_vector_observations
        self.vis_obs_size = brain.number_visual_observations
        tf.Variable(int(brain.vector_action_space_type == 'continuous'),
                    name='is_continuous_control', trainable=False, dtype=tf.int32)
        tf.Variable(self._version_number_, name='version_number', trainable=False, dtype=tf.int32)
        tf.Variable(self.m_size, name="memory_size", trainable=False, dtype=tf.int32)
        if brain.vector_action_space_type == 'continuous':
            tf.Variable(self.act_size[0], name="action_output_shape", trainable=False, dtype=tf.int32)
        else:
            tf.Variable(sum(self.act_size), name="action_output_shape", trainable=False, dtype=tf.int32)
//  [python]
def _normalize(self, A):
        n = tf.shape(A)[-1]
        I = tf.to_float(tf.eye(n))[tf.newaxis, :, :]
        return I
//  [python]
def _normalize(self, A):
        n = tf.shape(A)[-1]
        I = tf.cast(tf.eye(n), dtype=A.dtype)[tf.newaxis, :, :]
        return I
//  [python]
def _get_feature_vector_attention_mask(
        self, feature_vector_length: int, attention_mask: tf.Tensor, add_adapter=None
    ):
        non_padded_lengths = tf.math.cumsum(attention_mask, axis=-1)[:, -1]
        output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)
        output_lengths = tf.to_int32(output_lengths)
        batch_size = tf.shape(attention_mask)[0]
        # check device here
        attention_mask = tf.zeros(
            (batch_size, feature_vector_length), dtype=attention_mask.dtype, name="attention_mask"
        )  # these two operations makes sure that all values before the output lengths idxs are attended to
        ## check device
        attention_mask = tf.tensor_scatter_nd_update(
            attention_mask,
            indices=tf.stack([tf.range(batch_size), output_lengths - 1], axis=1),
            updates=tf.ones([batch_size], dtype=attention_mask.dtype),
        )
        attention_mask = tf.reverse(attention_mask, axis=[-1])
        attention_mask = tf.cumsum(attention_mask, axis=-1)
        attention_mask = tf.reverse(attention_mask, axis=[-1])
        attention_mask = tf.to_bool(attention_mask)
        return attention_mask
//  [python]
def _get_feature_vector_attention_mask(
        self, feature_vector_length: int, attention_mask: tf.Tensor, add_adapter=None
    ):
        non_padded_lengths = tf.math.cumsum(attention_mask, axis=-1)[:, -1]
        output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)
        output_lengths = tf.cast(output_lengths, tf.int32)
        batch_size = tf.shape(attention_mask)[0]
        # check device here
        attention_mask = tf.zeros(
            (batch_size, feature_vector_length), dtype=attention_mask.dtype, name="attention_mask"
        )  # these two operations makes sure that all values before the output lengths idxs are attended to
        ## check device
        attention_mask = tf.tensor_scatter_nd_update(
            attention_mask,
            indices=tf.stack([tf.range(batch_size), output_lengths - 1], axis=1),
            updates=tf.ones([batch_size], dtype=attention_mask.dtype),
        )
        attention_mask = tf.reverse(attention_mask, axis=[-1])
        attention_mask = tf.cumsum(attention_mask, axis=-1)
        attention_mask = tf.reverse(attention_mask, axis=[-1])
        attention_mask = tf.cast(attention_mask, tf.bool)
        return attention_mask
//  [python]
def _mask_hidden_states(self, hidden_states: tf.Tensor, mask_time_indices: tf.Tensor | None = None):
        """
        Masks extracted features along time axis and/or along feature axis according to
        [SpecAugment](https://arxiv.org/abs/1904.08779).
        """
        batch_size, sequence_length, hidden_size = shape_list(hidden_states)

        # `config.apply_spec_augment` can set masking to False
        if not getattr(self.config, "apply_spec_augment", True):
            return hidden_states

        if mask_time_indices is not None:
            # apply SpecAugment along time axis with given mask_time_indices
            hidden_states = tf.where(
                t.cast(mask_time_indices[:, :, tf.newaxis], tf.bool),
                self.masked_spec_embed[tf.newaxis, tf.newaxis, :],
                hidden_states,
            )

        elif self.config.mask_time_prob > 0:
            # generate indices & apply SpecAugment along time axis
            mask_time_indices = _compute_mask_indices(
                (batch_size, sequence_length),
                mask_prob=self.config.mask_time_prob,
                mask_length=self.config.mask_time_length,
                min_masks=2,
            )
            hidden_states = tf.where(
                t.cast(mask_time_indices[:, :, tf.newaxis], tf.bool),
                self.masked_spec_embed[tf.newaxis, tf.newaxis, :],
                hidden_states,
            )

        # apply SpecAugment along feature axis
        if self.config.mask_feature_prob > 0:
            mask_feature_indices = _compute_mask_indices(
                (batch_size, hidden_size),
                mask_prob=self.config.mask_feature_prob,
                mask_length=self.config.mask_feature_length,
            )
            hidden_states = tf.where(mask_feature_indices[:, tf.newaxis, :], hidden_states, 0)

        return hidden_states