jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __call__(self,
               inputs=None,
               encoder_mask=None,
               dense_inputs=None,
               training=False):
    """Applies Transformer model on the inputs.

    Args:
      inputs: input word ids. Optional if dense data are provided.
      encoder_mask: the encoder self-attention mask.
      dense_inputs: dense input data. Concat after the embedding if word ids are
        provided.
      training: whether it is training pass, affecting dropouts.

    Returns:
      output of a transformer encoder.
    """
    # Casts inputs to the dtype.
    if encoder_mask is not None:
      encoder_mask = tf.cast(encoder_mask, self.compute_dtype)
    cfg = self.config
    inputs_array = []
    if inputs is not None:
      inputs_array.append(
          self.input_embed(inputs, one_hot=cfg.one_hot_embedding))
    if dense_inputs is not None:
      inputs_array.append(dense_inputs)
    if not inputs_array:
      raise ValueError("At least one of inputs and dense_inputs must not be "
                       "None.")
    x = tf.concat(inputs_array, axis=1)
    tensor_shape = tf_utils.get_shape_list(x)
    tensor_shape[-2] = 1
    x = self.input_dropout(x, noise_shape=tensor_shape, training=training)
    if inputs is not None:
      input_length = tf_utils.get_shape_list(inputs)[1]
    else:
      input_length = 0

    attention_outputs = []
    for i in range(cfg.num_layers):
      position_bias = self.get_relpos_bias(input_length, dense_inputs, i)
      x = self.encoder_layers[i](
          x,
          attention_mask=encoder_mask,
          position_bias=position_bias,
          training=training)
      if self.config.return_attention_scores:
        x, attention_scores = x
        attention_outputs.append(attention_scores)

    encoded = self.output_norm(x)
    encoded = self.output_dropout(encoded, training=training)
    if self.config.return_attention_scores:
      return encoded, attention_outputs
    else:
      return encoded
//  [python]
def __init__(self,
               d_model: int,
               d_kv: int,
               num_heads: int,
               use_bias: bool = False,
               dropout_rate: Optional[float] = 0.0,
               rescale_query: bool = False,
               weight_initializer: Optional[Initializer] = None,
               bias_initializer: Optional[Initializer] = None,
               **kwargs):
    super().__init__(**kwargs)
    with self.name_scope:
      self.d_model = d_model
      self.d_kv = d_kv
      self.num_heads = num_heads
      self.rescale_query = rescale_query
      self.use_bias = use_bias

      if rescale_query or weight_initializer is None:
        query_w_init = weight_initializer
      else:
        init_std_rescaling = t.cast(self.d_kv, dtype=self.dtype)
        query_w_init = (
            lambda *args, **kwargs: (  # pylint: disable=g-long-lambda
                tf_utils.clone_initializer(weight_initializer)
                (*args, **kwargs) / init_std_rescaling))
      self.q = Linear3D(
          self.d_model,
          self.d_kv,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          w_init=query_w_init,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="q")
      self.k = Linear3D(
          self.d_model,
          self.d_kv,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          w_init=weight_initializer,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="k")
      self.v = Linear3D(
          self.d_model,
          self.d_kv,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          w_init=weight_initializer,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="v")
      self.o = Linear3D(
          self.d_kv,
          self.d_model,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          to_3d=False,
          w_init=weight_initializer,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="o")
      self.dropout = Dropout(dropout_rate)
//  [python]
def __init__(self,
               d_model: int,
               d_kv: int,
               num_heads: int,
               use_bias: bool = False,
               dropout_rate: Optional[float] = 0.0,
               rescale_query: bool = False,
               weight_initializer: Optional[Initializer] = None,
               bias_initializer: Optional[Initializer] = None,
               **kwargs):
    super().__init__(**kwargs)
    with self.name_scope:
      self.d_model = d_model
      self.d_kv = d_kv
      self.num_heads = num_heads
      self.rescale_query = rescale_query
      self.use_bias = use_bias

      if rescale_query or weight_initializer is None:
        query_w_init = weight_initializer
      else:
        init_std_rescaling = tf.math.sqrt(tf.cast(self.d_kv, dtype=self.dtype))
        query_w_init = (
            lambda *args, **kwargs: (  # pylint: disable=g-long-lambda
                tf_utils.clone_initializer(weight_initializer)
                (*args, **kwargs) / init_std_rescaling))
      self.q = Linear3D(
          self.d_model,
          self.d_kv,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          w_init=query_w_init,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="q")
      self.k = Linear3D(
          self.d_model,
          self.d_kv,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          w_init=weight_initializer,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="k")
      self.v = Linear3D(
          self.d_model,
          self.d_kv,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          w_init=weight_initializer,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="v")
      self.o = Linear3D(
          self.d_kv,
          self.d_model,
          num_heads=self.num_heads,
          use_bias=self.use_bias,
          to_3d=False,
          w_init=weight_initializer,
          b_init=bias_initializer,
          dtype=self.dtype,
          name="o")
      self.dropout = Dropout(dropout_rate)
//  [python]
def encode(self,
             encoder_input_tokens=None,
             encoder_segment_ids=None,
             encoder_dense_inputs=None,
             encoder_dense_segment_ids=None,
             training=False):
    eligible_position_array = []
    if encoder_input_tokens is not None:
      eligible_position_array.append(
          tf.to_float(tf.not_equal(encoder_input_tokens, 0)))
    if encoder_dense_inputs is not None:
      eligible_dense_positions = tf.to_float(
          tf.reduce_any(tf.not_equal(encoder_dense_inputs, 0), axis=-1))
      eligible_position_array.append(eligible_dense_positions)
    if not eligible_position_array:
      raise ValueError("At least one of encoder_input_tokens and"
                       " encoder_dense_inputs must be provided.")

    eligible_positions = tf.concat(eligible_position_array, axis=1)
    encoder_mask = make_attention_mask(
        eligible_positions, eligible_positions, dtype=tf.bool)

    encoder_segment_id_array = []
    if encoder_segment_ids is not None:
      encoder_segment_id_array.append(encoder_segment_ids)
    if encoder_dense_segment_ids is not None:
      encoder_segment_id_array.append(encoder_dense_segment_ids)
    if encoder_segment_id_array:
      encoder_segment_ids = tf.concat(encoder_segment_id_array, axis=1)
      segment_mask = make_attention_mask(
          encoder_segment_ids, encoder_segment_ids, tf.equal, dtype=tf.bool)
      encoder_mask = tf.math.logical_and(encoder_mask, segment_mask)
    encoder_mask = (1.0 - tf.to_float(encoder_mask)) * -1e9
    return self.encoder(
        encoder_input_tokens,
        encoder_mask,
        encoder_dense_inputs,
        training=training)
//  [python]
def encode(self,
             encoder_input_tokens=None,
             encoder_segment_ids=None,
             encoder_dense_inputs=None,
             encoder_dense_segment_ids=None,
             training=False):
    eligible_position_array = []
    if encoder_input_tokens is not None:
      eligible_position_array.append(
          tf.cast(tf.not_equal(encoder_input_tokens, 0), self.compute_dtype))
    if encoder_dense_inputs is not None:
      eligible_dense_positions = tf.cast(
          tf.reduce_any(tf.not_equal(encoder_dense_inputs, 0), axis=-1),
          self.compute_dtype)
      eligible_position_array.append(eligible_dense_positions)
    if not eligible_position_array:
      raise ValueError("At least one of encoder_input_tokens and"
                       " encoder_dense_inputs must be provided.")

    eligible_positions = tf.concat(eligible_position_array, axis=1)
    encoder_mask = make_attention_mask(
        eligible_positions, eligible_positions, dtype=tf.bool)

    encoder_segment_id_array = []
    if encoder_segment_ids is not None:
      encoder_segment_id_array.append(encoder_segment_ids)
    if encoder_dense_segment_ids is not None:
      encoder_segment_id_array.append(encoder_dense_segment_ids)
    if encoder_segment_id_array:
      encoder_segment_ids = tf.concat(encoder_segment_id_array, axis=1)
      segment_mask = make_attention_mask(
          encoder_segment_ids, encoder_segment_ids, tf.equal, dtype=tf.bool)
      encoder_mask = tf.math.logical_and(encoder_mask, segment_mask)
    encoder_mask = (1.0 - tf.cast(encoder_mask, self.compute_dtype)) * -1e9
    return self.encoder(
        encoder_input_tokens,
        encoder_mask,
        encoder_dense_inputs,
        training=training)
//  [python]
def dot_op(self, features):
        batch_size = tf.shape(features)[0]
        matrixdot = tf.matmul(features, features, transpose_b=True)
        feature_dim = matrixdot.shape[-1]

        ones_mat = tf.ones_like(matrixdot)
        lower_tri_mat = ones_mat - tf.linalg.band_part(ones_mat, 0, -1)
        lower_tri_mask = tf.to_float(lower_tri_mat) > 0
        result = tf.boolean_mask(matrixdot, lower_tri_mask)
        output_dim = feature_dim * (feature_dim - 1) // 2

        return tf.reshape(result, (batch_size, output_dim))
//  [python]
def dot_op(self, features):
        batch_size = tf.shape(features)[0]
        matrixdot = tf.matmul(features, features, transpose_b=True)
        feature_dim = matrixdot.shape[-1]

        ones_mat = tf.ones_like(matrixdot)
        lower_tri_mat = ones_mat - tf.linalg.band_part(ones_mat, 0, -1)
        lower_tri_mask = tf.cast(lower_tri_mat, tf.bool)
        result = tf.boolean_mask(matrixdot, lower_tri_mask)
        output_dim = feature_dim * (feature_dim - 1) // 2

        return tf.reshape(result, (batch_size, output_dim))
//  [python]
def centernet_proto_get(serialized):
  """Extracts the contents from a VoxelSample proto to tensors."""
  message_type = 'giotto_occluded_primitives.CenterNetSample'
  field_names = ['name', 'image_data', 'image_size', 'center2d', 'center3d',
                 'box_dims2d', 'box_dims3d', 'rotation', 'translation', 'rt',
                 'k']
  output_types = [tf.string, tf.string, tf.int32, tf.float32, tf.float32,
                  tf.float32, tf.float32, tf.float32, tf.float32, tf.float32,
                  tf.float32]
  _, values = decode_proto(serialized, message_type, field_names, output_types)
  filename = tf.squeeze(values[0])
  image_data = tf.squeeze(values[1])
  image_size = tf.squeeze(values[2])
  center2d = tf.squeeze(values[3])
  center3d = tf.squeeze(values[4])
  box_dims2d = tf.squeeze(values[5])
  box_dims3d = tf.squeeze(values[6])
  rotation = tf.squeeze(values[7])
  translation = tf.squeeze(values[8])
  rt = tf.squeeze(values[9])
  k = tf.squeeze(values[10])
  image = tf.image.convert_image_dtype(tf.io.decode_image(image_data, channels=3), dtype=tf.float32)
  return filename, \
         image, \
         image_size, \
         center2d, \
         center3d, \
         box_dims2d, \
         box_dims3d, \
         rotation, \
         translation, \
         rt, \
         k
//  [python]
def centernet_proto_get(serialized):
  """Extracts the contents from a VoxelSample proto to tensors."""
  message_type = 'giotto_occluded_primitives.CenterNetSample'
  field_names = ['name', 'image_data', 'image_size', 'center2d', 'center3d',
                 'box_dims2d', 'box_dims3d', 'rotation', 'translation', 'rt',
                 'k']
  output_types = [tf.string, tf.string, tf.int32, tf.float32, tf.float32,
                  tf.float32, tf.float32, tf.float32, tf.float32, tf.float32,
                  tf.float32]
  _, values = decode_proto(serialized, message_type, field_names, output_types)
  filename = tf.squeeze(values[0])
  image_data = tf.squeeze(values[1])
  image_size = tf.squeeze(values[2])
  center2d = tf.squeeze(values[3])
  center3d = tf.squeeze(values[4])
  box_dims2d = tf.squeeze(values[5])
  box_dims3d = tf.squeeze(values[6])
  rotation = tf.squeeze(values[7])
  translation = tf.squeeze(values[8])
  rt = tf.squeeze(values[9])
  k = tf.squeeze(values[10])
  image = tf.cast(tf.io.decode_image(image_data, channels=3), tf.float32)
  return filename, \
         image, \
         image_size, \
         center2d, \
         center3d, \
         box_dims2d, \
         box_dims3d, \
         rotation, \
         translation, \
         rt, \
         k
//  [python]
def __call__(self,
               decoder_input_tokens,
               encoded,
               decoder_mask=None,
               encoder_decoder_mask=None,
               decode=False,
               decode_position=None,
               cache=None,
               max_decode_len=None,
               training=False):
    """Applies Transformer model on the inputs.

    Args:
      decoder_input_tokens: the decoder input tokens.
      encoded: the encoder outputs.
      decoder_mask: the decoder self-attention mask.
      encoder_decoder_mask: the cross-attention mask.
      decode: Whether to perform autoregressive decoding.
      decode_position: integer, the position to decode.
      cache: The cache dictionary of key, value tensors.
      max_decode_len: An optional integer specifying the maximum decoding
        length. Note that this is only used for defining the relative position
        embedding parameters.
      training: Whether it is training pass, affecting dropouts.

    Returns:
      output of a transformer encoder including
      1. logits: Logits for each word in the vocab.
      2. raw_logits: Logits along the moded dimension.
      3. cache: Used for decoding in inference mode.
    """
    cfg = self.config
    # Casts inputs to the dtype.
    encoded = t.cast(encoded, self.compute_dtype)
    if decoder_mask is not None:
      decoder_mask = t.cast(decoder_mask, self.compute_dtype)
    if encoder_decoder_mask is not None:
      encoder_decoder_mask = t.cast(encoder_decoder_mask, self.compute_dtype)
    x = self.target_embed(decoder_input_tokens, one_hot=cfg.one_hot_embedding)
    tensor_shape = tf_utils.get_shape_list(x)
    tensor_shape[-2] = 1
    x = self.target_dropout(x, noise_shape=tensor_shape, training=training)

    for i in range(cfg.num_decoder_layers):
      if cache is not None:
        position_bias = self.get_relpos_bias(max_decode_len, i)
      else:
        input_length = tf_utils.get_shape_list(decoder_input_tokens)[1]
        position_bias = self.get_relpos_bias(input_length, i)

      if cache is None:
        x, _ = self.decoder_layers[i](
            x,
            encoder_hidden_states=encoded,
            attention_mask=decoder_mask,
            encoder_decoder_mask=encoder_decoder_mask,
            position_bias=position_bias,
            training=training)
      else:
        x, cache[i] = self.decoder_layers[i](
            x,
            encoder_hidden_states=encoded,
            attention_mask=decoder_mask,
            encoder_decoder_mask=encoder_decoder_mask,
            position_bias=position_bias,
            decode_position=decode_position,
            cache=cache[i],
            training=training)

    output = self.output_norm(x)
    tensor_shape = tf_utils.get_shape_list(output)
    tensor_shape[-2] = 1
    output = self.target_dropout(
        output, noise_shape=tensor_shape, training=training)
    if self.config.logits_via_embedding:
      logits = self.target_embed.attend(output)
      logits = logits / math.sqrt(cfg.d_model)
    else:
      logits = self.logits_dense(output)
    return dict(logits=logits, cache=cache, raw_logits=output)
//  [python]
def __call__(self,
               decoder_input_tokens,
               encoded,
               decoder_mask=None,
               encoder_decoder_mask=None,
               decode=False,
               decode_position=None,
               cache=None,
               max_decode_len=None,
               training=False):
    """Applies Transformer model on the inputs.

    Args:
      decoder_input_tokens: the decoder input tokens.
      encoded: the encoder outputs.
      decoder_mask: the decoder self-attention mask.
      encoder_decoder_mask: the cross-attention mask.
      decode: Whether to perform autoregressive decoding.
      decode_position: integer, the position to decode.
      cache: The cache dictionary of key, value tensors.
      max_decode_len: An optional integer specifying the maximum decoding
        length. Note that this is only used for defining the relative position
        embedding parameters.
      training: Whether it is training pass, affecting dropouts.

    Returns:
      output of a transformer encoder including
      1. logits: Logits for each word in the vocab.
      2. raw_logits: Logits along the moded dimension.
      3. cache: Used for decoding in inference mode.
    """
    cfg = self.config
    # Casts inputs to the dtype.
    encoded = tf.cast(encoded, self.compute_dtype)
    if decoder_mask is not None:
      decoder_mask = tf.cast(decoder_mask, self.compute_dtype)
    if encoder_decoder_mask is not None:
      encoder_decoder_mask = tf.cast(encoder_decoder_mask, self.compute_dtype)
    x = self.target_embed(decoder_input_tokens, one_hot=cfg.one_hot_embedding)
    tensor_shape = tf_utils.get_shape_list(x)
    tensor_shape[-2] = 1
    x = self.target_dropout(x, noise_shape=tensor_shape, training=training)

    for i in range(cfg.num_decoder_layers):
      if cache is not None:
        position_bias = self.get_relpos_bias(max_decode_len, i)
      else:
        input_length = tf_utils.get_shape_list(decoder_input_tokens)[1]
        position_bias = self.get_relpos_bias(input_length, i)

      if cache is None:
        x, _ = self.decoder_layers[i](
            x,
            encoder_hidden_states=encoded,
            attention_mask=decoder_mask,
            encoder_decoder_mask=encoder_decoder_mask,
            position_bias=position_bias,
            training=training)
      else:
        x, cache[i] = self.decoder_layers[i](
            x,
            encoder_hidden_states=encoded,
            attention_mask=decoder_mask,
            encoder_decoder_mask=encoder_decoder_mask,
            position_bias=position_bias,
            decode_position=decode_position,
            cache=cache[i],
            training=training)

    output = self.output_norm(x)
    tensor_shape = tf_utils.get_shape_list(output)
    tensor_shape[-2] = 1
    output = self.target_dropout(
        output, noise_shape=tensor_shape, training=training)
    if self.config.logits_via_embedding:
      logits = self.target_embed.attend(output)
      logits = logits / math.sqrt(cfg.d_model)
    else:
      logits = self.logits_dense(output)
    return dict(logits=logits, cache=cache, raw_logits=output)
//  [python]
def mask_logits(logits, mask):
        mask = t.cast(mask, t.float32)
        logits = logits * mask - (1 - mask) * 1e12
        return logits
//  [python]
def mask_logits(logits, mask):
        mask = tf.cast(mask, tf.float32)
        logits = logits * mask - (1 - mask) * 1e12
        return logits
//  [python]
def warmup_and_cosine_lr(init_lr, global_step, warmup_step, decay_steps, num_gpu, alpha=1e-6):
    def warmup(end_lr, global_step, warmup_step):
        start_lr = end_lr * 0.1
        global_step = t.cast(global_step, t.float32)
        return start_lr + (end_lr - start_lr) * global_step / warmup_step

    def cosine_lr(init_lr, global_step, decay_steps, num_gpu, alpha=0.0):
        return tf.train.cosine_decay(learning_rate=init_lr,
                                     global_step=global_step - warmup_step,
                                     decay_steps=decay_steps // num_gpu - warmup_step,
                                     alpha=alpha)

    return tf.cond(tf.less_equal(global_step, warmup_step),
                   true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                   false_fn=lambda: cosine_lr(init_lr, global_step, decay_steps, num_gpu, alpha))
//  [python]
def warmup_and_cosine_lr(init_lr, global_step, warmup_step, decay_steps, num_gpu, alpha=1e-6):
    def warmup(end_lr, global_step, warmup_step):
        start_lr = end_lr * 0.1
        global_step = tf.cast(global_step, tf.float32)
        return start_lr + (end_lr - start_lr) * global_step / warmup_step

    def cosine_lr(init_lr, global_step, decay_steps, num_gpu, alpha=0.0):
        return tf.train.cosine_decay(learning_rate=init_lr,
                                     global_step=global_step - warmup_step,
                                     decay_steps=decay_steps // num_gpu - warmup_step,
                                     alpha=alpha)

    return tf.cond(tf.less_equal(global_step, warmup_step),
                   true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                   false_fn=lambda: cosine_lr(init_lr, global_step, decay_steps, num_gpu, alpha))
//  [python]
def GRU(seq, seq_length=None, mask=None, name='text'):
        k = seq.get_shape().as_list()[-1]
        if seq_length is None and mask is not None:
            seq_length = tf.reduce_sum(tf.to_int32(mask), -1)
        with tf.variable_scope(f'GRU_{name}', reuse=tf.AUTO_REUSE):
            cell = tf.nn.rnn_cell.GRUCell(k)
            # ret: outputs, state
            outputs, state = tf.nn.dynamic_rnn(cell, seq, sequence_length=seq_length, dtype=tf.float32)
            return outputs, state
//  [python]
def GRU(seq, seq_length=None, mask=None, name='text'):
        k = seq.get_shape().as_list()[-1]
        if seq_length is None and mask is not None:
            seq_length = tf.reduce_sum(tf.cast(mask, tf.int32), -1)
        with tf.variable_scope(f'GRU_{name}', reuse=tf.AUTO_REUSE):
            cell = tf.nn.rnn_cell.GRUCell(k)
            # ret: outputs, state
            outputs, state = tf.nn.dynamic_rnn(cell, seq, sequence_length=seq_length, dtype=tf.float32)
            return outputs, state
//  [python]
def get_next_sentence_output(bert_config, input_tensor, labels):
  """Get loss and log probs for the next sentence prediction."""

  # Simple binary classification. Note that 0 is "next sentence" and 1 is
  # "random sentence". This weight matrix is not used after pre-training.
  with tf.variable_scope("cls/seq_relationship"):
    output_weights = tf.get_variable(
        "output_weights",
        shape=[2, bert_config.hidden_size],
        initializer=modeling.create_initializer(bert_config.initializer_range))
    output_bias = tf.get_variable(
        "output_bias", shape=[2], initializer=tf.zeros_initializer())

    logits = tf.matmul(tf.to_float(input_tensor), output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    log_probs = tf.nn.log_softmax(logits, axis=-1)
    labels = tf.reshape(labels, [-1])
    one_hot_labels = tf.one_hot(labels, depth=2, dtype=tf.float32)
    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
    loss = tf.reduce_mean(per_example_loss)
    return (loss, per_example_loss, log_probs)
//  [python]
def get_next_sentence_output(bert_config, input_tensor, labels):
  """Get loss and log probs for the next sentence prediction."""

  # Simple binary classification. Note that 0 is "next sentence" and 1 is
  # "random sentence". This weight matrix is not used after pre-training.
  with tf.variable_scope("cls/seq_relationship"):
    output_weights = tf.get_variable(
        "output_weights",
        shape=[2, bert_config.hidden_size],
        initializer=modeling.create_initializer(bert_config.initializer_range))
    output_bias = tf.get_variable(
        "output_bias", shape=[2], initializer=tf.zeros_initializer())

    logits = tf.matmul(tf.cast(input_tensor, tf.float32), output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    log_probs = tf.nn.log_softmax(logits, axis=-1)
    labels = tf.reshape(labels, [-1])
    one_hot_labels = tf.one_hot(labels, depth=2, dtype=tf.float32)
    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
    loss = tf.reduce_mean(per_example_loss)
    return (loss, per_example_loss, log_probs)
//  [python]
def _to_float(tensor):
  return t.cast(tensor, tf.float32)
//  [python]
def _to_float(tensor):
  return tf.cast(tensor, tf.float32)
//  [python]
def get_masked_lm_output(bert_config, input_tensor, output_weights, positions,
                         label_ids, label_weights):
  """Get loss and log probs for the masked LM."""
  input_tensor = gather_indexes(input_tensor, positions)

  with tf.variable_scope("cls/predictions"):
    # We apply one more non-linear transformation before the output layer.
    # This matrix is not used after pre-training.
    with tf.variable_scope("transform"):
      input_tensor = tf.layers.dense(
          input_tensor,
          units=bert_config.hidden_size,
          activation=modeling.get_activation(bert_config.hidden_act),
          kernel_initializer=modeling.create_initializer(
              bert_config.initializer_range))
      input_tensor = modeling.layer_norm(input_tensor)

    # The output weights are the same as the input embeddings, but there is
    # an output-only bias for each token.
    output_bias = tf.get_variable(
        "output_bias",
        shape=[bert_config.vocab_size],
        initializer=tf.zeros_initializer())
    logits = tf.matmul(tf.to_float(input_tensor), output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    label_ids = tf.reshape(label_ids, [-1])
    label_weights = tf.reshape(label_weights, [-1])

    one_hot_labels = tf.one_hot(
        label_ids, depth=bert_config.vocab_size, dtype=tf.float32)

    # The `positions` tensor might be zero-padded (if the sequence is too
    # short to have the maximum number of predictions). The `label_weights`
    # tensor has a value of 1.0 for every real prediction and 0.0 for the
    # padding predictions.
    per_example_loss = -tf.reduce_sum(log_probs * one_hot_labels, axis=[-1])
    numerator = tf.reduce_sum(label_weights * per_example_loss)
    denominator = tf.reduce_sum(label_weights) + 1e-5
    loss = numerator / denominator

  return (loss, per_example_loss, log_probs)
//  [python]
def get_masked_lm_output(bert_config, input_tensor, output_weights, positions,
                         label_ids, label_weights):
  """Get loss and log probs for the masked LM."""
  input_tensor = gather_indexes(input_tensor, positions)

  with tf.variable_scope("cls/predictions"):
    # We apply one more non-linear transformation before the output layer.
    # This matrix is not used after pre-training.
    with tf.variable_scope("transform"):
      input_tensor = tf.layers.dense(
          input_tensor,
          units=bert_config.hidden_size,
          activation=modeling.get_activation(bert_config.hidden_act),
          kernel_initializer=modeling.create_initializer(
              bert_config.initializer_range))
      input_tensor = modeling.layer_norm(input_tensor)

    # The output weights are the same as the input embeddings, but there is
    # an output-only bias for each token.
    output_bias = tf.get_variable(
        "output_bias",
        shape=[bert_config.vocab_size],
        initializer=tf.zeros_initializer())
    logits = tf.matmul(tf.cast(input_tensor, tf.float32), output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    label_ids = tf.reshape(label_ids, [-1])
    label_weights = tf.reshape(label_weights, [-1])

    one_hot_labels = tf.one_hot(
        label_ids, depth=bert_config.vocab_size, dtype=tf.float32)

    # The `positions` tensor might be zero-padded (if the sequence is too
    # short to have the maximum number of predictions). The `label_weights`
    # tensor has a value of 1.0 for every real prediction and 0.0 for the
    # padding predictions.
    per_example_loss = -tf.reduce_sum(log_probs * one_hot_labels, axis=[-1])
    numerator = tf.reduce_sum(label_weights * per_example_loss)
    denominator = tf.reduce_sum(label_weights) + 1e-5
    loss = numerator / denominator

  return (loss, per_example_loss, log_probs)
//  [python]
def Mean(seq, seq_length=None, mask=None, name=None):
        # seq: (None, L, k), seq_length: (None, ), mask: (None, L)
        # ret: (None, k)
        if seq_length is None and mask is None:
            with tf.variable_scope('Mean'):
                return tf.reduce_sum(seq, -2)

        with tf.variable_scope('MaskMean'):
            if mask is None:
                mask = tf.sequence_mask(seq_length, maxlen=tf.shape(seq)[1], dtype=tf.float32)
            mask = t.cast(mask, tf.float32)
            mask = tf.expand_dims(mask, -1)  # (None, L, 1)
            seq = seq * mask
            seq = tf.reduce_sum(seq, -2)  # (None, k)
            seq = seq / (tf.reduce_sum(mask, -2) + eps)
        return seq
//  [python]
def Mean(seq, seq_length=None, mask=None, name=None):
        # seq: (None, L, k), seq_length: (None, ), mask: (None, L)
        # ret: (None, k)
        if seq_length is None and mask is None:
            with tf.variable_scope('Mean'):
                return tf.reduce_sum(seq, -2)

        with tf.variable_scope('MaskMean'):
            if mask is None:
                mask = tf.sequence_mask(seq_length, maxlen=tf.shape(seq)[1], dtype=tf.float32)
            mask = tf.cast(mask, tf.float32)
            mask = tf.expand_dims(mask, -1)  # (None, L, 1)
            seq = seq * mask
            seq = tf.reduce_sum(seq, -2)  # (None, k)
            seq = seq / (tf.reduce_sum(mask, -2) + eps)
        return seq
//  [python]
def get_metric_v(self, x, predict_v):
        # [BS,]
        true_item = x.ans
        true_item_a1 = tf.expand_dims(true_item, -1)
        # [BS, M], [BS, 1]
        eq = tf.cast(tf.equal(predict_v.top_items, true_item_a1), tf.int32)
        # [BS,]
        m = tf.reduce_max(eq, -1)
        idx = tf.cast(tf.argmax(eq, -1), tf.int32)
        rank = idx + m - 1
        ndcg = tf.log(2.0) * tf.cast(m, tf.float32) / tf.log(2.0 + tf.cast(idx, tf.float32))
        hit_rate = tf.cast(m, tf.float32)
        ret = Object(
            ndcg=ndcg,
            hit_rate=hit_rate,
            user=x.user,
            true_item=true_item,
            phase=x.phase,
            top_items=predict_v.top_items,
            top_scores=predict_v.top_scores,
            rank=rank,
            q_ts=x.q_ts,
        )
        return ret
//  [python]
def tf_draw_seed(seed):
    # TF ops only accept int32/64 seeds but our base seed is uint32.
    return t.cast(draw_seed(seed), dtype="int32")
//  [python]
def tf_draw_seed(seed):
    # TF ops only accept int32/64 seeds but our base seed is uint32.
    return tf.cast(draw_seed(seed), dtype="int32")
//  [python]
def randint(shape, minval, maxval, dtype="int32", seed=None):
    intemediate_dtype = dtype
    if standardize_dtype(dtype) not in ["int32", "int64"]:
        intemediate_dtype = "int64"
    seed = tf_draw_seed(seed)
    output = tf.random.stateless_uniform(
        shape=shape,
        minval=minval,
        maxval=maxval,
        dtype=intemediate_dtype,
        seed=seed,
    )
    return t.cast(output, dtype)
//  [python]
def randint(shape, minval, maxval, dtype="int32", seed=None):
    intemediate_dtype = dtype
    if standardize_dtype(dtype) not in ["int32", "int64"]:
        intemediate_dtype = "int64"
    seed = tf_draw_seed(seed)
    output = tf.random.stateless_uniform(
        shape=shape,
        minval=minval,
        maxval=maxval,
        dtype=intemediate_dtype,
        seed=seed,
    )
    return tf.cast(output, dtype)
//  [python]
def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
    dtype = dtype or floatx()
    seed = tf_draw_seed(seed)
    return tf.random.stateless_uniform(
        shape=shape,
        minval=t.cast(minval, dtype),
        maxval=t.cast(maxval, dtype),
        dtype=dtype,
        seed=seed,
    )
//  [python]
def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
    dtype = dtype or floatx()
    seed = tf_draw_seed(seed)
    return tf.random.stateless_uniform(
        shape=shape,
        minval=tf.cast(minval, dtype),
        maxval=tf.cast(maxval, dtype),
        dtype=dtype,
        seed=seed,
    )
//  [python]
def decode_bytes(serialized):
  """Extracts the contents from a VoxelSample proto to tensors."""
  message_type = 'giotto_occluded_primitives.CenterNetSample'
  field_names = ['name', 'image_data', 'image_size', 'center2d', 'center3d',
                 'box_dims2d', 'box_dims3d', 'rotation', 'translation', 'rt',
                 'k']
  output_types = [tf.string, tf.string, tf.int32, tf.float32, tf.float32,
                  tf.float32, tf.float32, tf.float32, tf.float32, tf.float32,
                  tf.float32]
  _, tensors = decode_proto(serialized, message_type, field_names, output_types)
  tensors = [tf.squeeze(t) for t in tensors]
  tensor_dict = dict(zip(field_names, tensors))
  tensor_dict['image'] = tf.cast(
      tf.io.decode_image(tensor_dict['image_data'], channels=3),
      tf.float32)
  tensor_dict['image'].set_shape([None, None, 3])
  tensor_dict['original_image_spatial_shape'] = \
    tf.shape(tensor_dict['image'])[:2]

  # Compute ground truth boxes. From (center2d,size2d) --> (ymin,xmin,ymax,xmax)
  box_min = tensor_dict['center2d'] - tensor_dict['box_dims2d']/2.0
  box_max = tensor_dict['center2d'] + tensor_dict['box_dims2d']/2.0
  box_min = tf.divide(box_min,
                      tf.cast(tensor_dict['original_image_spatial_shape'],
                              tf.float32))
  box_max = tf.divide(box_max,
                      tf.cast(tensor_dict['original_image_spatial_shape'],
                              tf.float32))
  tensor_dict['groundtruth_boxes'] = tf.reshape(
      tf.concat([box_min[1], box_min[0], box_max[1], box_max[0]], axis=0),
      [1, 4])

  tensor_dict['groundtruth_valid_classes'] = [1]
  tensor_dict['num_boxes'] = 1  # needs to be a scalar, no tensor
  tensor_dict['per_class_weight'] = [1]

  tensor_dict['box_dims3d'] = tf.reshape(tensor_dict['box_dims3d'], [1, 3])
  tensor_dict['size_offset3d'] = tensor_dict['box_dims3d'] - 0.5

  # Compute the dot on the ground plane used as prior for 3d pose (batched)
  image_size = tensor_dict['image_size']
  ray_2d = tf.cast(image_size, tf.float32) * [0.5, 3.0/4.0, 1/image_size[-1]]
  ray_2d = tf.reshape(ray_2d, [3, 1])
  k = tf.reshape(tensor_dict['k'], [3, 3])
  k_inv = tf.linalg.inv(k)
  rt = tf.reshape(tensor_dict['rt'], [3, 4])
  r = tf.gather(rt, [0, 1, 2], axis=1)
  t = tf.gather(rt, [3], axis=1)
  r_inv = tf.transpose(r, [1, 0])
  t_inv = tf.matmul(r_inv, t) * -1
  ray = tf.matmul(r_inv, tf.matmul(k_inv, ray_2d))
  l = -t_inv[-1] / ray[-1]  # determine lambda
  dot = tf.expand_dims(l, -1) * ray + t_inv
  tensor_dict['dot'] = tf.transpose(dot)
  tensor_dict['center3d'] = tf.reshape(tensor_dict['center3d'], [1, 3])
  tensor_dict['dot_offset3d'] = tensor_dict['center3d'] - tensor_dict['dot']

  # Rotation
  rad = tensor_dict['rotation'] * -1 * math.pi/180
  cos = tf.reshape(tf.cos(rad), [1, 1])
  sin = tf.reshape(tf.sin(rad), [1, 1])
  tensor_dict['rotation'] = tf.concat([cos - 1,
                                       sin - 0,
                                       -sin- 0,
                                       cos - 1], axis=1)
  return tensor_dict
//  [python]
def gamma(shape, alpha, dtype=None, seed=None):
    dtype = dtype or floatx()
    seed = tf_draw_seed(seed)
    # TODO: `tf.random.stateless_gamma` doesn't support bfloat16
    intemediate_dtype = dtype
    if standardize_dtype(dtype) == "bfloat16":
        intemediate_dtype = "float32"
    return t.cast(
        tf.random.stateless_gamma(
            shape,
            alpha=alpha,
            dtype=intemediate_dtype,
            seed=seed,
        ),
        dtype,
    )
//  [python]
def gamma(shape, alpha, dtype=None, seed=None):
    dtype = dtype or floatx()
    seed = tf_draw_seed(seed)
    # TODO: `tf.random.stateless_gamma` doesn't support bfloat16
    intemediate_dtype = dtype
    if standardize_dtype(dtype) == "bfloat16":
        intemediate_dtype = "float32"
    return tf.cast(
        tf.random.stateless_gamma(
            shape,
            alpha=alpha,
            dtype=intemediate_dtype,
            seed=seed,
        ),
        dtype,
    )
//  [python]
def binomial(shape, counts, probabilities, dtype=None, seed=None):
    dtype = dtype or floatx()
    seed = tf_draw_seed(seed)
    # TODO: `tf.random.stateless_binomial` doesn't support bfloat16
    intemediate_dtype = dtype
    if standardize_dtype(dtype) == "bfloat16":
        intemediate_dtype = "float32"
    return tf.cast(
        tf.random.stateless_binomial(
            shape=shape,
            seed=seed,
            counts=counts,
            probs=probabilities,
            output_dtype=intemediate_dtype,
        ),
        dtype,
    )
//  [python]
def get_predict_v(self, x, score):
        # [N]
        item_pop = t.cast(self.data.item_pop, t.float32)
        item_pop_log = t.log(item_pop + np.e)

        item_deg_self = t.cast(t.gather(self.data.item_deg_self_per_phase, x.phase), t.float32)
        item_pop_self_log = t.log(item_deg_self + np.e)

        if args.mode_pop == 'log':
            score = score * args.alpha_pop_base + score / item_pop_log
        elif args.mode_pop == 'log_mdeg':
            score = score * item_pop_self_log / item_pop_log
        elif args.mode_pop == 'log_mdeg_only':
            score = score * item_pop_self_log
        elif args.mode_pop == 'linear':
            item_pop = t.cast(self.data.item_pop, t.float32) + 1.0
            score = score * args.alpha_pop_base + score / item_pop
        elif args.mode_pop == 'log_md':
            item_pop_self_log = t.log(item_deg_self + np.e + 10.0)
            score = score * item_pop_self_log / item_pop_log

        if args.mode_rare in {'log', 'linear', 'base'}:
            if args.mode_rare == 'log':
                rare_weight_pop = 1.0 / t.log(t.cast(self.data.item_pop, t.float32) + np.e)
            elif args.mode_rare == 'linear':
                rare_weight_pop = 1.0 / t.cast(self.data.item_pop, t.float32)
            elif args.mode_rare == 'base':
                rare_weight_pop = 0.0
            else:
                raise Exception
            # [N]
            rare_weight = float(args.alpha_rare)
            rare_weight = rare_weight + rare_weight_pop
            rare_weight *= float(args.alpha_rare_mul)

            is_rare = self.is_rare(x)
            score = t.where(is_rare, score * rare_weight + float(args.alpha_rare_base), score)

        score = UTILS.mask_logits(score, x.score_mask)

        t.summary.histogram('score', score)
        top_items, top_scores = self.topk_idx(score, x)

        if args.dump_all:
            self.tmp_vars.update(all_scores=score, item_seq=x.seq, ts_seq=x.ts, q_ts=x.q_ts)
            ret = Object(user=x.user, phase=x.phase, top_items=top_items, top_scores=top_scores)
            ret.update(**self.tmp_vars)
            return ret

        return Object(user=x.user, phase=x.phase, top_items=top_items, top_scores=top_scores)
//  [python]
def get_predict_v(self, x, score):
        # [N]
        item_pop = tf.cast(self.data.item_pop, tf.float32)
        item_pop_log = tf.log(item_pop + np.e)

        item_deg_self = tf.cast(tf.gather(self.data.item_deg_self_per_phase, x.phase), tf.float32)
        item_pop_self_log = tf.log(item_deg_self + np.e)

        if args.mode_pop == 'log':
            score = score * args.alpha_pop_base + score / item_pop_log
        elif args.mode_pop == 'log_mdeg':
            score = score * item_pop_self_log / item_pop_log
        elif args.mode_pop == 'log_mdeg_only':
            score = score * item_pop_self_log
        elif args.mode_pop == 'linear':
            item_pop = tf.cast(self.data.item_pop, tf.float32) + 1.0
            score = score * args.alpha_pop_base + score / item_pop
        elif args.mode_pop == 'log_md':
            item_pop_self_log = tf.log(item_deg_self + np.e + 10.0)
            score = score * item_pop_self_log / item_pop_log

        if args.mode_rare in {'log', 'linear', 'base'}:
            if args.mode_rare == 'log':
                rare_weight_pop = 1.0 / tf.log(tf.cast(self.data.item_pop, tf.float32) + np.e)
            elif args.mode_rare == 'linear':
                rare_weight_pop = 1.0 / tf.cast(self.data.item_pop, tf.float32)
            elif args.mode_rare == 'base':
                rare_weight_pop = 0.0
            else:
                raise Exception
            # [N]
            rare_weight = float(args.alpha_rare)
            rare_weight = rare_weight + rare_weight_pop
            rare_weight *= float(args.alpha_rare_mul)

            is_rare = self.is_rare(x)
            score = tf.where(is_rare, score * rare_weight + float(args.alpha_rare_base), score)

        score = UTILS.mask_logits(score, x.score_mask)

        tf.summary.histogram('score', score)
        top_items, top_scores = self.topk_idx(score, x)

        if args.dump_all:
            self.tmp_vars.update(all_scores=score, item_seq=x.seq, ts_seq=x.ts, q_ts=x.q_ts)
            ret = Object(user=x.user, phase=x.phase, top_items=top_items, top_scores=top_scores)
            ret.update(**self.tmp_vars)
            return ret

        return Object(user=x.user, phase=x.phase, top_items=top_items, top_scores=top_scores)
//  [python]
def __init__(self, config: MobileViTConfig, hidden_size: int, **kwargs) -> None:
        super().__init__(**kwargs)

        if hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                f"The hidden size {hidden_size,} is not a multiple of the number of attention "
                f"heads {config.num_attention_heads}."
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        scale = tf.cast(self.attention_head_size, dtype=tf.float32)
        self.scale = tf.math.sqrt(scale)

        self.query = tf.keras.layers.Dense(self.all_head_size, use_bias=config.qkv_bias, name="query")
        self.key = tf.keras.layers.Dense(self.all_head_size, use_bias=config.qkv_bias, name="key")
        self.value = tf.keras.layers.Dense(self.all_head_size, use_bias=config.qkv_bias, name="value")

        self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob)
//  [python]
def read_single_example_and_decode(self, filename_queue):

        # tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
        # reader = tf.TFRecordReader(options=tfrecord_options)
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)

        features = tf.parse_single_example(
            serialized=serialized_example,
            features={
                'img_name': tf.FixedLenFeature([], tf.string),
                'img_height': tf.FixedLenFeature([], tf.int64),
                'img_width': tf.FixedLenFeature([], tf.int64),
                'img': tf.FixedLenFeature([], tf.string),
                'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
                'num_objects': tf.FixedLenFeature([], tf.int64)
            }
        )
        img_name = features['img_name']
        img_height = tf.to_int32(features['img_height'])
        img_width = tf.to_int32(features['img_width'])
        img = tf.decode_raw(features['img'], tf.uint8)

        img = tf.reshape(img, shape=[img_height, img_width, 3])

        gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
        gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 9])

        num_objects = tf.to_int32(features['num_objects'])
        return img_name, img, gtboxes_and_label, num_objects
//  [python]
def read_single_example_and_decode(self, filename_queue):

        # tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
        # reader = tf.TFRecordReader(options=tfrecord_options)
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)

        features = tf.parse_single_example(
            serialized=serialized_example,
            features={
                'img_name': tf.FixedLenFeature([], tf.string),
                'img_height': tf.FixedLenFeature([], tf.int64),
                'img_width': tf.FixedLenFeature([], tf.int64),
                'img': tf.FixedLenFeature([], tf.string),
                'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
                'num_objects': tf.FixedLenFeature([], tf.int64)
            }
        )
        img_name = features['img_name']
        img_height = tf.cast(features['img_height'], tf.int32)
        img_width = tf.cast(features['img_width'], tf.int32)
        img = tf.decode_raw(features['img'], tf.uint8)

        img = tf.reshape(img, shape=[img_height, img_width, 3])

        gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
        gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 9])

        num_objects = tf.cast(features['num_objects'], tf.int32)
        return img_name, img, gtboxes_and_label, num_objects
//  [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 Fourier 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)
    return tf.math.real(self.fourier_transform(query))
//  [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 Fourier 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)
    return tf.math.real(self.fourier_transform(query))
//  [python]
def transform_targets(y_train, anchors, anchor_masks, size):
    y_outs = []
    grid_size = size // 32

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

    y_train = t.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)