jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def transform_targets(y_train, anchors, anchor_masks, size):
    y_outs = []
    grid_size = size // 32

    # calculate anchor index for true boxes
    anchors = tf.cast(anchors, tf.float32)
    anchor_area = anchors[..., 0] * anchors[..., 1]
    box_wh = y_train[..., 2:4] - y_train[..., 0:2]
    box_wh = tf.tile(tf.expand_dims(box_wh, -2),
                     (1, 1, tf.shape(anchors)[0], 1))
    box_area = box_wh[..., 0] * box_wh[..., 1]
    intersection = tf.minimum(box_wh[..., 0], anchors[..., 0]) * \
        tf.minimum(box_wh[..., 1], anchors[..., 1])
    iou = intersection / (box_area + anchor_area - intersection)
    anchor_idx = tf.cast(tf.argmax(iou, axis=-1), tf.float32)
    anchor_idx = tf.expand_dims(anchor_idx, axis=-1)

    y_train = tf.concat([y_train, anchor_idx], axis=-1)

    for anchor_idxs in anchor_masks:
        y_outs.append(transform_targets_for_output(
            y_train, grid_size, anchor_idxs))
        grid_size *= 2

    return tuple(y_outs)
//  [python]
def unfolding(self, features: tf.Tensor) -> Tuple[tf.Tensor, Dict]:
        patch_width, patch_height = self.patch_width, self.patch_height
        patch_area = tf.to_int32(patch_width * patch_height)

        batch_size = tf.shape(features)[0]
        orig_height = tf.shape(features)[1]
        orig_width = tf.shape(features)[2]
        channels = tf.shape(features)[3]

        new_height = tf.to_int32(tf.math.ceil(orig_height / patch_height) * patch_height)
        new_width = tf.to_int32(tf.math.ceil(orig_width / patch_width) * patch_width)

        interpolate = new_width != orig_width or new_height != orig_height
        if interpolate:
            # Note: Padding can be done, but then it needs to be handled in attention function.
            features = tf.image.resize(features, size=(new_height, new_width), method="bilinear")

        # number of patches along width and height
        num_patch_width = new_width // patch_width
        num_patch_height = new_height // patch_height
        num_patches = num_patch_height * num_patch_width

        # convert from shape (batch_size, orig_height, orig_width, channels)
        # to the shape (batch_size * patch_area, num_patches, channels)
        features = tf.transpose(features, [0, 3, 1, 2])
        patches = tf.reshape(
            features, (batch_size * channels * num_patch_height, patch_height, num_patch_width, patch_width)
        )
        patches = tf.transpose(patches, [0, 2, 1, 3])
        patches = tf.reshape(patches, (batch_size, channels, num_patches, patch_area))
        patches = tf.transpose(patches, [0, 3, 2, 1])
        patches = tf.reshape(patches, (batch_size * patch_area, num_patches, channels))

        info_dict = {
            "orig_size": (orig_height, orig_width),
            "batch_size": batch_size,
            "channels": channels,
            "interpolate": interpolate,
            "num_patches": num_patches,
            "num_patches_width": num_patch_width,
            "num_patches_height": num_patch_height,
        }
        return patches, info_dict
//  [python]
def unfolding(self, features: tf.Tensor) -> Tuple[tf.Tensor, Dict]:
        patch_width, patch_height = self.patch_width, self.patch_height
        patch_area = tf.cast(patch_width * patch_height, "int32")

        batch_size = tf.shape(features)[0]
        orig_height = tf.shape(features)[1]
        orig_width = tf.shape(features)[2]
        channels = tf.shape(features)[3]

        new_height = tf.cast(tf.math.ceil(orig_height / patch_height) * patch_height, "int32")
        new_width = tf.cast(tf.math.ceil(orig_width / patch_width) * patch_width, "int32")

        interpolate = new_width != orig_width or new_height != orig_height
        if interpolate:
            # Note: Padding can be done, but then it needs to be handled in attention function.
            features = tf.image.resize(features, size=(new_height, new_width), method="bilinear")

        # number of patches along width and height
        num_patch_width = new_width // patch_width
        num_patch_height = new_height // patch_height
        num_patches = num_patch_height * num_patch_width

        # convert from shape (batch_size, orig_height, orig_width, channels)
        # to the shape (batch_size * patch_area, num_patches, channels)
        features = tf.transpose(features, [0, 3, 1, 2])
        patches = tf.reshape(
            features, (batch_size * channels * num_patch_height, patch_height, num_patch_width, patch_width)
        )
        patches = tf.transpose(patches, [0, 2, 1, 3])
        patches = tf.reshape(patches, (batch_size, channels, num_patches, patch_area))
        patches = tf.transpose(patches, [0, 3, 2, 1])
        patches = tf.reshape(patches, (batch_size * patch_area, num_patches, channels))

        info_dict = {
            "orig_size": (orig_height, orig_width),
            "batch_size": batch_size,
            "channels": channels,
            "interpolate": interpolate,
            "num_patches": num_patches,
            "num_patches_width": num_patch_width,
            "num_patches_height": num_patch_height,
        }
        return patches, info_dict
//  [python]
def call(self, query: tf.Tensor, value: tf.Tensor, **kwargs) -> tf.Tensor:
    """Applies layer to `query`.

    Args:
      query: Batch of input embeddings, typically of shape <float>[batch_size,
        max_seq_length, hidden_dim].
      value: Unused. Included to match attention layer API.
      **kwargs: Optional arguments to catch unused attention keyword arguments.

    Returns:
      Real part of discrete Hartley Transform of `query` inputs with shape
        <float32>[batch_size, max_seq_length, hidden_dim].
    """
    del value  # Ignored by encoder-only mixing layers
    query = tf.to_float(query)
    frequencies = self.fourier_transform(query)
    return tf.math.real(frequencies) - tf.math.imag(frequencies)
//  [python]
def call(self, query: tf.Tensor, value: tf.Tensor, **kwargs) -> tf.Tensor:
    """Applies layer to `query`.

    Args:
      query: Batch of input embeddings, typically of shape <float>[batch_size,
        max_seq_length, hidden_dim].
      value: Unused. Included to match attention layer API.
      **kwargs: Optional arguments to catch unused attention keyword arguments.

    Returns:
      Real part of discrete Hartley Transform of `query` inputs with shape
        <float32>[batch_size, max_seq_length, hidden_dim].
    """
    del value  # Ignored by encoder-only mixing layers
    query = tf.cast(query, tf.complex64)
    frequencies = self.fourier_transform(query)
    return tf.math.real(frequencies) - tf.math.imag(frequencies)
//  [python]
def _resize_detection_masks(args):
  detection_boxes, detection_masks, image_shape = args
  detection_masks_reframed = ops.reframe_box_masks_to_image_masks(
      detection_masks, detection_boxes, image_shape[0], image_shape[1])
  return tf.to_int64(tf.greater(detection_masks_reframed, 0.5))
//  [python]
def _resize_detection_masks(args):
  detection_boxes, detection_masks, image_shape = args
  detection_masks_reframed = ops.reframe_box_masks_to_image_masks(
      detection_masks, detection_boxes, image_shape[0], image_shape[1])
  return tf.cast(tf.greater(detection_masks_reframed, 0.5), tf.uint8)
//  [python]
def transform_targets_for_output(y_true, grid_size, anchor_idxs):
    # y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
    N = tf.shape(y_true)[0]

    # y_true_out: (N, grid, grid, anchors, [x1, y1, x2, y2, obj, class])
    y_true_out = tf.zeros(
        (N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))

    anchor_idxs = tf.to_int32(anchor_idxs)

    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    for i in tf.range(N):
        for j in tf.range(tf.shape(y_true)[1]):
            if tf.equal(y_true[i][j][2], 0):
                continue
            anchor_eq = tf.equal(
                anchor_idxs, tf.to_int32(y_true[i][j][5]))

            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][0:4]
                box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2

                anchor_idx = tf.to_int32(tf.where(anchor_eq))
                grid_xy = tf.to_int32(box_xy // (1/grid_size))

                # grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
                indexes = indexes.write(
                    idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
                updates = updates.write(
                    idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
                idx += 1

    # tf.print(indexes.stack())
    # tf.print(updates.stack())

    return tf.tensor_scatter_nd_update(
        y_true_out, indexes.stack(), updates.stack())
//  [python]
def transform_targets_for_output(y_true, grid_size, anchor_idxs):
    # y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
    N = tf.shape(y_true)[0]

    # y_true_out: (N, grid, grid, anchors, [x1, y1, x2, y2, obj, class])
    y_true_out = tf.zeros(
        (N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))

    anchor_idxs = tf.cast(anchor_idxs, tf.int32)

    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    for i in tf.range(N):
        for j in tf.range(tf.shape(y_true)[1]):
            if tf.equal(y_true[i][j][2], 0):
                continue
            anchor_eq = tf.equal(
                anchor_idxs, tf.cast(y_true[i][j][5], tf.int32))

            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][0:4]
                box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2

                anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)
                grid_xy = tf.cast(box_xy // (1/grid_size), tf.int32)

                # grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
                indexes = indexes.write(
                    idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
                updates = updates.write(
                    idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
                idx += 1

    # tf.print(indexes.stack())
    # tf.print(updates.stack())

    return tf.tensor_scatter_nd_update(
        y_true_out, indexes.stack(), updates.stack())
//  [python]
def read_and_prepocess_single_img(self, filename_queue, shortside_len, is_training):

        img_name, img, gtboxes_and_label, num_objects = self.read_single_example_and_decode(filename_queue)

        img = t.cast(img, t.float32)

        if is_training:

            if self.cfgs.RGB2GRAY:
                # img, gtboxes_and_label = image_preprocess.aspect_ratio_jittering(img, gtboxes_and_label)
                img = self.image_preprocess.random_rgb2gray(img_tensor=img, gtboxes_and_label=gtboxes_and_label)

            if self.cfgs.IMG_ROTATE:
                img, gtboxes_and_label = self.image_preprocess.random_rotate_img(img_tensor=img,
                                                                                 gtboxes_and_label=gtboxes_and_label)

            img, gtboxes_and_label, img_h, img_w = self.image_preprocess.short_side_resize(img_tensor=img,
                                                                                           gtboxes_and_label=gtboxes_and_label,
                                                                                           target_shortside_len=shortside_len,
                                                                                           length_limitation=self.cfgs.IMG_MAX_LENGTH)

            if self.cfgs.HORIZONTAL_FLIP:
                img, gtboxes_and_label = self.image_preprocess.random_flip_left_right(img_tensor=img,
                                                                                      gtboxes_and_label=gtboxes_and_label)
            if self.cfgs.VERTICAL_FLIP:
                img, gtboxes_and_label = self.image_preprocess.random_flip_up_down(img_tensor=img,
                                                                                   gtboxes_and_label=gtboxes_and_label)

        else:
            img, gtboxes_and_label, img_h, img_w = self.image_preprocess.short_side_resize(img_tensor=img,
                                                                                           gtboxes_and_label=gtboxes_and_label,
                                                                                           target_shortside_len=shortside_len,
                                                                                           length_limitation=self.cfgs.IMG_MAX_LENGTH)
        pretrain_zoo = PretrainModelZoo()
        if self.cfgs.NET_NAME in pretrain_zoo.pth_zoo or self.cfgs.NET_NAME in pretrain_zoo.mxnet_zoo:
            img = img / 255 - tf.constant([[self.cfgs.PIXEL_MEAN_]])
        else:
            img = img - tf.constant([[self.cfgs.PIXEL_MEAN]])  # sub pixel mean at last
        return img_name, img, gtboxes_and_label, num_objects, img_h, img_w
//  [python]
def read_and_prepocess_single_img(self, filename_queue, shortside_len, is_training):

        img_name, img, gtboxes_and_label, num_objects = self.read_single_example_and_decode(filename_queue)

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

        if is_training:

            if self.cfgs.RGB2GRAY:
                # img, gtboxes_and_label = image_preprocess.aspect_ratio_jittering(img, gtboxes_and_label)
                img = self.image_preprocess.random_rgb2gray(img_tensor=img, gtboxes_and_label=gtboxes_and_label)

            if self.cfgs.IMG_ROTATE:
                img, gtboxes_and_label = self.image_preprocess.random_rotate_img(img_tensor=img,
                                                                                 gtboxes_and_label=gtboxes_and_label)

            img, gtboxes_and_label, img_h, img_w = self.image_preprocess.short_side_resize(img_tensor=img,
                                                                                           gtboxes_and_label=gtboxes_and_label,
                                                                                           target_shortside_len=shortside_len,
                                                                                           length_limitation=self.cfgs.IMG_MAX_LENGTH)

            if self.cfgs.HORIZONTAL_FLIP:
                img, gtboxes_and_label = self.image_preprocess.random_flip_left_right(img_tensor=img,
                                                                                      gtboxes_and_label=gtboxes_and_label)
            if self.cfgs.VERTICAL_FLIP:
                img, gtboxes_and_label = self.image_preprocess.random_flip_up_down(img_tensor=img,
                                                                                   gtboxes_and_label=gtboxes_and_label)

        else:
            img, gtboxes_and_label, img_h, img_w = self.image_preprocess.short_side_resize(img_tensor=img,
                                                                                           gtboxes_and_label=gtboxes_and_label,
                                                                                           target_shortside_len=shortside_len,
                                                                                           length_limitation=self.cfgs.IMG_MAX_LENGTH)
        pretrain_zoo = PretrainModelZoo()
        if self.cfgs.NET_NAME in pretrain_zoo.pth_zoo or self.cfgs.NET_NAME in pretrain_zoo.mxnet_zoo:
            img = img / 255 - tf.constant([[self.cfgs.PIXEL_MEAN_]])
        else:
            img = img - tf.constant([[self.cfgs.PIXEL_MEAN]])  # sub pixel mean at last
        return img_name, img, gtboxes_and_label, num_objects, img_h, img_w
//  [python]
def _resize_groundtruth_masks(args):
  mask, image_shape = args
  mask = tf.expand_dims(mask, 3)
  mask = tf.image.resize_images(
      mask,
      image_shape,
      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
      align_corners=True)
  return tf.to_int64(tf.squeeze(mask, 3))
//  [python]
def _resize_groundtruth_masks(args):
  mask, image_shape = args
  mask = tf.expand_dims(mask, 3)
  mask = tf.image.resize_images(
      mask,
      image_shape,
      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
      align_corners=True)
  return tf.cast(tf.squeeze(mask, 3), tf.uint8)
//  [python]
def _pick_fourier_transform(
    use_fft: bool, max_seq_length: int,
    hidden_dim: int) -> Callable[[tf.Tensor], tf.Tensor]:
  """Returns FFT or DFT Fourier Transform implementation.

  On TPUs, we recommend using the Discrete Fourier Transform (DFT) matrix
  (use_fft=False), except for very long sequence lengths. On GPUs and CPUs, the
  Fast Fourier Transform (use_fft=True) is generally optimal for all sequence
  lengths.

  Note: When using the FFT it is recommended to use a sequence length that is a
  power of 2.

  Args:
    use_fft: If True, return FFT. Otherwise, return DFT matrix.
    max_seq_length: Maximum sequence length of inputs. Only used if
      use_fft=False.
    hidden_dim: Size of hidden dimension of inputs. Only used if use_fft=False.

  Returns:
    Fourier Transform.
  """
  if use_fft:
    return tf.signal.fft2d
  else:
    dft_mat_seq = linalg.dft(max_seq_length).astype(np.complex64)
    dft_mat_hidden = linalg.dft(hidden_dim).astype(np.complex64)

    def two_dim_matmul(x: tf.Tensor, matrix_dim_one: tf.Tensor,
                       matrix_dim_two: tf.Tensor) -> tf.Tensor:
      """Applies 2D matrix multiplication to input tensors of rank >= 2."""
      return tf.einsum("...ij,jk,ni->...nk", tf.complex(x.real, x.imag),
                       matrix_dim_two, matrix_dim_one)

    return functools.partial(
        two_dim_matmul,
        matrix_dim_one=dft_mat_seq,
        matrix_dim_two=dft_mat_hidden)
//  [python]
def _pick_fourier_transform(
    use_fft: bool, max_seq_length: int,
    hidden_dim: int) -> Callable[[tf.Tensor], tf.Tensor]:
  """Returns FFT or DFT Fourier Transform implementation.

  On TPUs, we recommend using the Discrete Fourier Transform (DFT) matrix
  (use_fft=False), except for very long sequence lengths. On GPUs and CPUs, the
  Fast Fourier Transform (use_fft=True) is generally optimal for all sequence
  lengths.

  Note: When using the FFT it is recommended to use a sequence length that is a
  power of 2.

  Args:
    use_fft: If True, return FFT. Otherwise, return DFT matrix.
    max_seq_length: Maximum sequence length of inputs. Only used if
      use_fft=False.
    hidden_dim: Size of hidden dimension of inputs. Only used if use_fft=False.

  Returns:
    Fourier Transform.
  """
  if use_fft:
    return tf.signal.fft2d
  else:
    dft_mat_seq = linalg.dft(max_seq_length).astype(np.complex64)
    dft_mat_hidden = linalg.dft(hidden_dim).astype(np.complex64)

    def two_dim_matmul(x: tf.Tensor, matrix_dim_one: tf.Tensor,
                       matrix_dim_two: tf.Tensor) -> tf.Tensor:
      """Applies 2D matrix multiplication to input tensors of rank >= 2."""
      return tf.einsum("...ij,jk,ni->...nk", tf.cast(x, tf.complex64),
                       matrix_dim_two, matrix_dim_one)

    return functools.partial(
        two_dim_matmul,
        matrix_dim_one=dft_mat_seq,
        matrix_dim_two=dft_mat_hidden)
//  [python]
def length(sequence):
    """
    This function return the sequence length of each x in the batch.
    :param sequence: the batch sequence of shape [batch_size, num_steps, feature_size]
    :return length: A tensor of shape [batch_size]
    """
    used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
    seq_length = tf.reduce_sum(used, 1)
    seq_length = tf.to_int32(seq_length)
    return seq_length
//  [python]
def length(sequence):
    """
    This function return the sequence length of each x in the batch.
    :param sequence: the batch sequence of shape [batch_size, num_steps, feature_size]
    :return length: A tensor of shape [batch_size]
    """
    used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
    seq_length = tf.reduce_sum(used, 1)
    seq_length = tf.cast(seq_length, tf.int32)
    return seq_length
//  [python]
def process_batch_input(input_batch):
  input_batch = tf.reshape(input_batch, [tf.shape(input_batch)[0], -1])
  input_batch = tf.to_float(input_batch)
  return input_batch
//  [python]
def process_batch_input(input_batch):
  input_batch = tf.reshape(input_batch, [tf.shape(input_batch)[0], -1])
  input_batch = tf.cast(input_batch, tf.float32)
  return input_batch
//  [python]
def call(
        self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None, training: bool = False
    ) -> tf.Tensor:
        embeddings = self.patch_embeddings(pixel_values)
        batch_size, seq_length, _ = shape_list(embeddings)

        if bool_masked_pos is not None:
            mask_tokens = tf.tile(self.mask_token, [batch_size, seq_length, 1])
            # replace the masked visual tokens by mask_tokens
            mask = tf.expand_dims(bool_masked_pos, axis=-1)
            mask = tf.to_float(mask) 
            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask

        cls_tokens = tf.repeat(self.cls_token, repeats=batch_size, axis=0)
        distillation_tokens = tf.repeat(self.distillation_token, repeats=batch_size, axis=0)
        embeddings = tf.concat((cls_tokens, distillation_tokens, embeddings), axis=1)
        embeddings = embeddings + self.position_embeddings
        embeddings = self.dropout(embeddings, training=training)
        return embeddings
//  [python]
def call(
        self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None, training: bool = False
    ) -> tf.Tensor:
        embeddings = self.patch_embeddings(pixel_values)
        batch_size, seq_length, _ = shape_list(embeddings)

        if bool_masked_pos is not None:
            mask_tokens = tf.tile(self.mask_token, [batch_size, seq_length, 1])
            # replace the masked visual tokens by mask_tokens
            mask = tf.expand_dims(bool_masked_pos, axis=-1)
            mask = tf.cast(mask, dtype=mask_tokens.dtype)
            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask

        cls_tokens = tf.repeat(self.cls_token, repeats=batch_size, axis=0)
        distillation_tokens = tf.repeat(self.distillation_token, repeats=batch_size, axis=0)
        embeddings = tf.concat((cls_tokens, distillation_tokens, embeddings), axis=1)
        embeddings = embeddings + self.position_embeddings
        embeddings = self.dropout(embeddings, training=training)
        return embeddings
//  [python]
def _attn(self, q, k, v, attention_mask, head_mask, output_attentions, training=False):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        if self.scale:
            dk = tf.to_float(shape_list(k)[-1])  # scale attention_scores
            w = w / tf.math.sqrt(dk)

        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = tf.to_float(self.causal_attention_mask(nd, ns))
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w * b - 1e4 * (1 - b)

        if attention_mask is not None:
            # Apply the attention mask
            attention_mask = tf.to_float(attention_mask)
            w = w + attention_mask

        w = stable_softmax(w, axis=-1)
        w = self.attn_dropout(w, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            w = w * head_mask

        outputs = [tf.matmul(w, v)]
        if output_attentions:
            outputs.append(w)
        return outputs
//  [python]
def _attn(self, q, k, v, attention_mask, head_mask, output_attentions, training=False):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        if self.scale:
            dk = tf.cast(shape_list(k)[-1], dtype=w.dtype)  # scale attention_scores
            w = w / tf.math.sqrt(dk)

        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = tf.cast(self.causal_attention_mask(nd, ns), dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w * b - 1e4 * (1 - b)

        if attention_mask is not None:
            # Apply the attention mask
            attention_mask = tf.cast(attention_mask, dtype=w.dtype)
            w = w + attention_mask

        w = stable_softmax(w, axis=-1)
        w = self.attn_dropout(w, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            w = w * head_mask

        outputs = [tf.matmul(w, v)]
        if output_attentions:
            outputs.append(w)
        return outputs
//  [python]
def indices_to_dense_vector(indices,
                            size,
                            indices_value=1.,
                            default_value=0,
                            dtype=tf.float32):
  """Creates dense vector with indices set to specific value and rest to zeros.

  This function exists because it is unclear if it is safe to use
    tf.sparse_to_dense(indices, [size], 1, validate_indices=False)
  with indices which are not ordered.
  This function accepts a dynamic size (e.g. tf.shape(tensor)[0])

  Args:
    indices: 1d Tensor with integer indices which are to be set to
      indices_values.
    size: scalar with size (integer) of output Tensor.
    indices_value: values of elements specified by indices in the output vector
    default_value: values of other elements in the output vector.
    dtype: data type.

  Returns:
    dense 1D Tensor of shape [size] with indices set to indices_values and the
      rest set to default_value.
  """
  size = t.cast(size, dtype=t.int32)
  zeros = tf.ones([size], dtype=dtype) * default_value
  values = tf.ones_like(indices, dtype=dtype) * indices_value

  return tf.dynamic_stitch(
      [tf.range(size), t.cast(indices, dtype=t.int32)], [zeros, values])
//  [python]
def indices_to_dense_vector(indices,
                            size,
                            indices_value=1.,
                            default_value=0,
                            dtype=tf.float32):
  """Creates dense vector with indices set to specific value and rest to zeros.

  This function exists because it is unclear if it is safe to use
    tf.sparse_to_dense(indices, [size], 1, validate_indices=False)
  with indices which are not ordered.
  This function accepts a dynamic size (e.g. tf.shape(tensor)[0])

  Args:
    indices: 1d Tensor with integer indices which are to be set to
      indices_values.
    size: scalar with size (integer) of output Tensor.
    indices_value: values of elements specified by indices in the output vector
    default_value: values of other elements in the output vector.
    dtype: data type.

  Returns:
    dense 1D Tensor of shape [size] with indices set to indices_values and the
      rest set to default_value.
  """
  size = tf.cast(size, dtype=tf.int32)
  zeros = tf.ones([size], dtype=dtype) * default_value
  values = tf.ones_like(indices, dtype=dtype) * indices_value

  return tf.dynamic_stitch(
      [tf.range(size), tf.cast(indices, dtype=tf.int32)], [zeros, values])
//  [python]
def construct_combination_matrix(input_rankings, rankings_per_img):
        def batch_fn(batch_dim):
            def instance_fn(batch_segment):
                ranking_segments = t.cast(batch_segment[:, 1], dtype=t.int32)
                ranking_segments_counts = tf.math.segment_sum(tf.ones_like(ranking_segments), ranking_segments)
                segment_range = tf.range(tf.constant(0, dtype=tf.int32),
                                         tf.cast(tf.shape(ranking_segments_counts)[0], dtype=tf.int32), dtype=tf.int32)

                def inner_fn(range_ids):
                    def inner_inner_fn(ranking_segment_counts, range_id):
                        n_skip_elements = tf.reduce_sum(ranking_segment_counts[:range_id])
                        repeat_elements = tf.reduce_sum(ranking_segment_counts[range_id:], keepdims=True)

                        def construct_mesh(num_repeats):
                            a = tf.tile(tf.expand_dims(tf.constant([0, 1]), axis=-1), [1, num_repeats[0]])

                            arguments = tf.unstack(a, axis=1)
                            mesh_grid = tf.meshgrid(*arguments, indexing='ij')
                            return tf.reshape(tf.stack(mesh_grid, axis=-1), (-1, num_repeats[0]))[1:]

                        combs = construct_mesh(repeat_elements)
                        return tf.cond(tf.equal(n_skip_elements, 0), lambda: combs, lambda: tf.concat(
                            [tf.zeros([combs.shape[0], n_skip_elements], dtype=combs.dtype), combs], axis=1))

                    def py_fun_mesh(loc_x, loc_y):
                        return tf.py_function(inner_inner_fn, [loc_x, loc_y], tf.int32)

                    inner_result = py_fun_mesh(ranking_segments_counts, range_ids)

                    return tf.RaggedTensor.from_tensor(inner_result)

                return tf.map_fn(inner_fn, segment_range, fn_output_signature=tf.RaggedTensorSpec(shape=[None, None],
                                                                                                  dtype=tf.int32))

            return tf.map_fn(instance_fn, batch_dim, fn_output_signature=tf.RaggedTensorSpec(shape=[None, None, None],
                                                                                             dtype=tf.int32))

        return tf.map_fn(batch_fn, input_rankings,
                         fn_output_signature=tf.RaggedTensorSpec(shape=[rankings_per_img, None, None, None],
                                                                 dtype=tf.int32))
//  [python]
def construct_combination_matrix(input_rankings, rankings_per_img):
        def batch_fn(batch_dim):
            def instance_fn(batch_segment):
                ranking_segments = tf.cast(batch_segment[:, 1], dtype=tf.int32)
                ranking_segments_counts = tf.math.segment_sum(tf.ones_like(ranking_segments), ranking_segments)
                segment_range = tf.range(tf.constant(0, dtype=tf.int32),
                                         tf.cast(tf.shape(ranking_segments_counts)[0], dtype=tf.int32), dtype=tf.int32)

                def inner_fn(range_ids):
                    def inner_inner_fn(ranking_segment_counts, range_id):
                        n_skip_elements = tf.reduce_sum(ranking_segment_counts[:range_id])
                        repeat_elements = tf.reduce_sum(ranking_segment_counts[range_id:], keepdims=True)

                        def construct_mesh(num_repeats):
                            a = tf.tile(tf.expand_dims(tf.constant([0, 1]), axis=-1), [1, num_repeats[0]])

                            arguments = tf.unstack(a, axis=1)
                            mesh_grid = tf.meshgrid(*arguments, indexing='ij')
                            return tf.reshape(tf.stack(mesh_grid, axis=-1), (-1, num_repeats[0]))[1:]

                        combs = construct_mesh(repeat_elements)
                        return tf.cond(tf.equal(n_skip_elements, 0), lambda: combs, lambda: tf.concat(
                            [tf.zeros([combs.shape[0], n_skip_elements], dtype=combs.dtype), combs], axis=1))

                    def py_fun_mesh(loc_x, loc_y):
                        return tf.py_function(inner_inner_fn, [loc_x, loc_y], tf.int32)

                    inner_result = py_fun_mesh(ranking_segments_counts, range_ids)

                    return tf.RaggedTensor.from_tensor(inner_result)

                return tf.map_fn(inner_fn, segment_range, fn_output_signature=tf.RaggedTensorSpec(shape=[None, None],
                                                                                                  dtype=tf.int32))

            return tf.map_fn(instance_fn, batch_dim, fn_output_signature=tf.RaggedTensorSpec(shape=[None, None, None],
                                                                                             dtype=tf.int32))

        return tf.map_fn(batch_fn, input_rankings,
                         fn_output_signature=tf.RaggedTensorSpec(shape=[rankings_per_img, None, None, None],
                                                                 dtype=tf.int32))
//  [python]
def train_step(self, data: tuple[InputBatch, tf.Tensor, tf.Tensor]):
    features, labels, weights = data
    with tf.GradientTape() as tape:
      predictions = self.call(features)
      labels = tf.cast(labels, tf.float32)
      loss_vector = self._encoder.loss_vector(
          predictions, labels, weights=weights, features=features)
    self.optimizer.minimize(
        loss=loss_vector,
        var_list=self._encoder.trainable_variables,
        tape=tape)
    return {}
//  [python]
def _sliced_recall(labels: pd.DataFrame,
                   preds: tf.Tensor,
                   k: int,
                   buckets: Sequence[_Bucket]) -> Sequence[float]:
  """Returns the recall per bucket.

  Args:
    labels: a DataFrame of the labels, containing the columns "uid" and "sid".
    preds: a Tensor of predictions, of shape [num_users, num_items].
    k: computes the Recall@k.
    buckets: a list of lists of sid, used to partition the items into buckets.
  """
  labels = labels.sort_values(["uid", "sid"])
  num_items = max(np.max(labels["sid"]), np.max(preds)) + 1
  num_users = preds.shape[0]
  preds = preds[:, :k]
  recalls = []
  for bucket in buckets:
    l = labels[labels["sid"].isin(bucket)]
    sp_l = tf.SparseTensor(
        indices=l[["uid", "sid"]],
        values=tf.to_int32(l["sid"]),
        dense_shape=[num_users, num_items])
    true_pos = tf.sets.size(tf.sets.intersection(preds, sp_l))
    all_pos = tf.minimum(k, tf.sets.size(sp_l))
    recall = true_pos/all_pos
    recalls.append(np.nanmean(recall.numpy()))
  return recalls
//  [python]
def _sliced_recall(labels: pd.DataFrame,
                   preds: tf.Tensor,
                   k: int,
                   buckets: Sequence[_Bucket]) -> Sequence[float]:
  """Returns the recall per bucket.

  Args:
    labels: a DataFrame of the labels, containing the columns "uid" and "sid".
    preds: a Tensor of predictions, of shape [num_users, num_items].
    k: computes the Recall@k.
    buckets: a list of lists of sid, used to partition the items into buckets.
  """
  labels = labels.sort_values(["uid", "sid"])
  num_items = max(np.max(labels["sid"]), np.max(preds)) + 1
  num_users = preds.shape[0]
  preds = preds[:, :k]
  recalls = []
  for bucket in buckets:
    l = labels[labels["sid"].isin(bucket)]
    sp_l = tf.SparseTensor(
        indices=l[["uid", "sid"]],
        values=tf.cast(l["sid"], tf.int32),
        dense_shape=[num_users, num_items])
    true_pos = tf.sets.size(tf.sets.intersection(preds, sp_l))
    all_pos = tf.minimum(k, tf.sets.size(sp_l))
    recall = true_pos/all_pos
    recalls.append(np.nanmean(recall.numpy()))
  return recalls
//  [python]
def _recall(labels: tf.SparseTensor,
            preds: tf.Tensor,
            k: int,
            per_user: bool = False) -> tf.Tensor:
  """Returns the recall at k given predictions and sparse labels.

  Args:
    labels: a SparseTensor of the positive ids.
    preds: a dense tensor of the predicted scores.
    k: compute Recall@k.
    per_user: when True, returns a vector of recalls, one per user. Otherwise,
      returns a scalar of the overall recall (in this case every positive will
      have an equal weight, so users with more positives contribute more to the
      metric).
  """
  top_k = tf.to_int64(tf.math.top_k(preds, k)[1])
  true_pos = tf.sets.size(tf.sets.intersection(top_k, labels))
  all_pos = tf.minimum(k, tf.sets.size(labels))
  if per_user:
    return true_pos/all_pos
  def nan_mean(xs):
    mask = tf.reshape(tf.logical_not(tf.math.is_nan(xs)), [-1])
    return tf.reduce_mean(tf.boolean_mask(xs, mask))
  return nan_mean(true_pos/all_pos)
//  [python]
def _recall(labels: tf.SparseTensor,
            preds: tf.Tensor,
            k: int,
            per_user: bool = False) -> tf.Tensor:
  """Returns the recall at k given predictions and sparse labels.

  Args:
    labels: a SparseTensor of the positive ids.
    preds: a dense tensor of the predicted scores.
    k: compute Recall@k.
    per_user: when True, returns a vector of recalls, one per user. Otherwise,
      returns a scalar of the overall recall (in this case every positive will
      have an equal weight, so users with more positives contribute more to the
      metric).
  """
  top_k = tf.cast(tf.math.top_k(preds, k)[1], tf.int64)
  true_pos = tf.sets.size(tf.sets.intersection(top_k, labels))
  all_pos = tf.minimum(k, tf.sets.size(labels))
  if per_user:
    return true_pos/all_pos
  def nan_mean(xs):
    mask = tf.reshape(tf.logical_not(tf.math.is_nan(xs)), [-1])
    return tf.reduce_mean(tf.boolean_mask(xs, mask))
  return nan_mean(true_pos/all_pos)
//  [python]
def argmax(self, logits):
    """Samples the most likely class label given the logits.

    Args:
      logits: Un-normalized logits with shape (num_samples, num_classes)

    Returns:
      predictions: Predicted class labels, has shape (num_samples,)
    """
    return tf.to_int32(tf.argmax(tf.nn.softmax(logits), 1))
//  [python]
def argmax(self, logits):
    """Samples the most likely class label given the logits.

    Args:
      logits: Un-normalized logits with shape (num_samples, num_classes)

    Returns:
      predictions: Predicted class labels, has shape (num_samples,)
    """
    return tf.cast(tf.argmax(tf.nn.softmax(logits), 1), tf.int32)
//  [python]
def __init__(self,
               name: str,
               model: "_TwoTowerModel",
               matrix: SparseMatrix,
               eval_batch_size: Optional[int],
               recall_positions: Sequence[int]):
    """Computes eval metrics from a dataset.

    Compute user embeddings by projecting, using the history (only uses frequent
    items for projections).
    Note: some users may have 0 frequent items in their history, and will
    have a 0 embedding, i.e. their predictions are all 0.
    Evaluation includes all items, frequent and rare.

    Args:
      name: the name of the evaluation task
      model: a TwoTowerModel
      matrix: a SparseMatrix (train, validation, or test)
      eval_batch_size: used in wals_projection.
      recall_positions: a list of recall positions
    """
    self.name = name
    self.model = model
    self.num_items = model.num_items
    self.num_frequent_items = model.num_frequent_items
    self.matrix = matrix
    self.recall_positions = recall_positions
    # Batch the data for wals_projection.
    self.do_project = False
    if matrix.sp_history is not None:
      self.do_project = True
      self.history_ds = matrix.sp_history.batch_ls(eval_batch_size)
    # Pre-compute user-means for use in the metrics.
    self.user_avg_ratings = None
    if matrix.history is not None and not model.binarize:
      means = matrix.history.groupby("uid").mean()
      user_avg_ratings = np.zeros(matrix.num_users)
      user_avg_ratings[means.index] = means["rating"]
      self.user_avg_ratings = t.cast(user_avg_ratings, t.float32)
//  [python]
def __init__(self,
               name: str,
               model: "_TwoTowerModel",
               matrix: SparseMatrix,
               eval_batch_size: Optional[int],
               recall_positions: Sequence[int]):
    """Computes eval metrics from a dataset.

    Compute user embeddings by projecting, using the history (only uses frequent
    items for projections).
    Note: some users may have 0 frequent items in their history, and will
    have a 0 embedding, i.e. their predictions are all 0.
    Evaluation includes all items, frequent and rare.

    Args:
      name: the name of the evaluation task
      model: a TwoTowerModel
      matrix: a SparseMatrix (train, validation, or test)
      eval_batch_size: used in wals_projection.
      recall_positions: a list of recall positions
    """
    self.name = name
    self.model = model
    self.num_items = model.num_items
    self.num_frequent_items = model.num_frequent_items
    self.matrix = matrix
    self.recall_positions = recall_positions
    # Batch the data for wals_projection.
    self.do_project = False
    if matrix.sp_history is not None:
      self.do_project = True
      self.history_ds = matrix.sp_history.batch_ls(eval_batch_size)
    # Pre-compute user-means for use in the metrics.
    self.user_avg_ratings = None
    if matrix.history is not None and not model.binarize:
      means = matrix.history.groupby("uid").mean()
      user_avg_ratings = np.zeros(matrix.num_users)
      user_avg_ratings[means.index] = means["rating"]
      self.user_avg_ratings = tf.cast(user_avg_ratings, tf.float32)
//  [python]
def objective(self, params, data, labels):
    """Computes the softmax cross entropy.

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      loss: Softmax cross entropy loss averaged over the samples in the batch

    Raises:
      ValueError: If the objective is to be computed over >2 classes, because
        this operation is broken in tensorflow at the moment.
    """
    # Forward pass.
    logits = self.inference(params, data)

    # Compute the loss.
    l2reg = [tf.reduce_sum(param ** 2) for param in params]
    if int(logits.get_shape()[1]) == 2:
      labels = tf.to_float(labels)
      losses = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=labels, logits=logits[:, 0])
    else:
      raise ValueError("Unable to compute softmax cross entropy for more than"
                       " 2 classes.")

    return tf.reduce_mean(losses) + tf.reduce_mean(l2reg) * FLAGS.l2_reg_scale
//  [python]
def objective(self, params, data, labels):
    """Computes the softmax cross entropy.

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      loss: Softmax cross entropy loss averaged over the samples in the batch

    Raises:
      ValueError: If the objective is to be computed over >2 classes, because
        this operation is broken in tensorflow at the moment.
    """
    # Forward pass.
    logits = self.inference(params, data)

    # Compute the loss.
    l2reg = [tf.reduce_sum(param ** 2) for param in params]
    if int(logits.get_shape()[1]) == 2:
      labels = tf.cast(labels, tf.float32)
      losses = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=labels, logits=logits[:, 0])
    else:
      raise ValueError("Unable to compute softmax cross entropy for more than"
                       " 2 classes.")

    return tf.reduce_mean(losses) + tf.reduce_mean(l2reg) * FLAGS.l2_reg_scale
//  [python]
def inference(self, params, data):
    all_embeddings, softmax_weights, softmax_bias = params
    embeddings = tf.nn.embedding_lookup(all_embeddings, tf.to_int32(data))
    embeddings = tf.reduce_sum(embeddings, 1)
    return tf.matmul(embeddings, softmax_weights) + softmax_bias
//  [python]
def inference(self, params, data):
    all_embeddings, softmax_weights, softmax_bias = params
    embeddings = tf.nn.embedding_lookup(all_embeddings, tf.cast(data, tf.int32))
    embeddings = tf.reduce_sum(embeddings, 1)
    return tf.matmul(embeddings, softmax_weights) + softmax_bias
//  [python]
def accuracy(self, params, data, labels):
    """Computes the accuracy (fraction of correct classifications).

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      accuracy: Fraction of correct classifications across the batch
    """
    predictions = self.argmax(self.inference(params, data))
    return tf.contrib.metrics.accuracy(predictions, tf.cast(labels, tf.int32))
//  [python]
def accuracy(self, params, data, labels):
    """Computes the accuracy (fraction of correct classifications).

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      accuracy: Fraction of correct classifications across the batch
    """
    predictions = self.argmax(self.activation(self.inference(params, data)))
    return tf.contrib.metrics.accuracy(predictions, tf.cast(labels, tf.int32))
//  [python]
def inference(self, params, data):
    all_embeddings, softmax_weights, softmax_bias = params
    num_ids = tf.shape(data)[1]
    one_hot_embeddings = tf.one_hot(tf.to_int32(data), self.n_classes)
    one_hot_embeddings = tf.reshape(one_hot_embeddings, [-1, self.n_classes])
    embeddings = tf.matmul(one_hot_embeddings, all_embeddings)
    embeddings = tf.reshape(embeddings, [-1, num_ids, self.n_features])
    embeddings = tf.reduce_sum(embeddings, 1)
    return tf.matmul(embeddings, softmax_weights) + softmax_bias
//  [python]
def inference(self, params, data):
    all_embeddings, softmax_weights, softmax_bias = params
    num_ids = tf.shape(data)[1]
    one_hot_embeddings = tf.one_hot(tf.cast(data, tf.int32), self.n_classes)
    one_hot_embeddings = tf.reshape(one_hot_embeddings, [-1, self.n_classes])
    embeddings = tf.matmul(one_hot_embeddings, all_embeddings)
    embeddings = tf.reshape(embeddings, [-1, num_ids, self.n_features])
    embeddings = tf.reduce_sum(embeddings, 1)
    return tf.matmul(embeddings, softmax_weights) + softmax_bias
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    label_ids = labels[self.label_field]
    if self.task_config.num_classes == 1:
      loss = tf_keras.losses.mean_squared_error(label_ids, model_outputs)
    else:
      loss = tf_keras.losses.sparse_categorical_crossentropy(
          label_ids, tf.cast(model_outputs, tf.float32), from_logits=True
      )

    if aux_losses:
      loss += tf.add_n(aux_losses)
    return tf_utils.safe_mean(loss)
//  [python]
def mean(self):
    return t.cast(self._dist.mean(), self._dtype)
//  [python]
def mean(self):
    return tf.cast(self._dist.mean(), self._dtype)
//  [python]
def mode(self):
    return t.cast(self._dist.mode(), self._dtype)
//  [python]
def mode(self):
    return tf.cast(self._dist.mode(), self._dtype)
//  [python]
def entropy(self):
    return t.cast(self._dist.entropy(), self._dtype)
//  [python]
def entropy(self):
    return tf.cast(self._dist.entropy(), self._dtype)
//  [python]
def sample(self, *args, **kwargs):
    return tf.cast(self._dist.sample(*args, **kwargs), self._dtype)
//  [python]
def ae_loss_for_single_img(self,net,indexs,nr):
        indexs = indexs[:nr]
        fnr = t.cast(nr,t.float32)
        with t.name_scope("pull_loss"):
            net = t.reshape(net,[-1])
            old_shape = btf.combined_static_and_dynamic_shape(indexs)
            ids = t.gather(net,t.nn.relu(t.reshape(indexs,[-1])))
            ids = t.reshape(ids,old_shape)
            weights = t.cast(t.greater_equal(indexs,0),t.float32)
            mean = t.reduce_sum(ids*weights,axis=-1,keepdims=True)/t.maximum(t.reduce_sum(weights,axis=-1,keepdims=True),1e-8)
            pull_loss = t.square(ids-mean)*weights
            pull_loss = t.reduce_sum(pull_loss)/fnr

        with t.name_scope("push_loss"):
            range = t.range(nr)
            X,Y = t.meshgrid(range,range)
            X = t.reshape(X,[-1])
            Y = t.reshape(Y,[-1])
            mask = t.not_equal(X,Y)
            ftotal_nr = t.reduce_sum(t.cast(mask,t.float32))+1e-8
            X = t.boolean_mask(X,mask)
            Y = t.boolean_mask(Y,mask)
            loss = -t.square(t.gather(mean,X)-t.gather(mean,Y))
            loss = t.exp(loss)
            loss = t.reduce_sum(loss)
            loss = loss/ftotal_nr
            push_loss = loss*0.5

        return pull_loss,push_loss