jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def build_embed(inputs, features, index2words, keep_prob, is_training, create_embedding_lookup=True):
    embeddings = []
    word_embeddings = None
    for idx, (values, feature, index2word) in enumerate(zip(inputs, features, index2words)):
        if requires_vocab(feature):
            with tf.variable_scope("embedding_%d" % (idx,)):
                if create_embedding_lookup:
                    W, embedding = embedding_lookup(values,
                                                    dim=feature["dimension"],
                                                    size=len(index2word),
                                                    dtype=tf.float32,
                                                    mask_negative=True)
                    if feature["type"] == "word":
                        word_embeddings = (W, embedding, values)
                else:
                    embedding = tf.placeholder(tf.float32,
                                               feature_shape(feature) + [feature["dimension"]])
                tf.add_to_collection(EMBEDDED_INDICES, embedding)
                embeddings.append(embedding)
        elif requires_character_convolution(feature):
            embeddings.append(character_convolution(values, feature))
        elif feature["type"] == "bio":
            # b if different from previous tag
            start_is_b = tf.greater(values[:1], -1)
            b = tf.logical_and(tf.greater(values, -1), tf.concat([start_is_b, tf.not_equal(values[:-1], values[1:])], axis=0))
            i = tf.logical_and(tf.greater(values, -1), tf.concat([tf.zeros_like(start_is_b), tf.equal(values[:-1], values[1:])], axis=0))
            # O-tag is -1
            o = tf.less(values, 0)
            embeddings.append(tf.cast(tf.stack([b, i, o], axis=-1, name="bool_BIO"), tf.float32, name="BIO"))
        else:
            embeddings.append(tf.expand_dims(values, 2))
    return maybe_dropout(concat(embeddings, axis=2), keep_prob, is_training), word_embeddings
//  [python]
def soft_nn_loss(
    query_labels: IntTensor,
    query_embeddings: FloatTensor,
    key_labels: IntTensor,
    key_embeddings: FloatTensor,
    distance: Distance,
    temperature: float,
    remove_diagonal: bool = True,
) -> Any:
    """Computes the soft nearest neighbors loss.

    Args:
        query_labels: labels associated with the query embed.
        query_embeddings: Embedded query examples.
        key_labels: labels associated with the key embed.
        key_embeddings: Embedded key examples.
        distance: Which distance function to use to compute the pairwise.
        temperature: Controls relative importance given to the pair of points.
        remove_diagonal: Bool. If True, will set diagonal to False in positive pair mask

    Returns:
        loss: loss value for the current batch.
    """

    batch_size = tf.size(query_labels)
    eps = t.cast(1e-9, dtype=query_embeddings.dtype)

    pairwise_dist = distance(query_embeddings, key_embeddings)
    pairwise_dist = pairwise_dist / temperature
    negexpd = tf.math.exp(-pairwise_dist)

    # Mask out diagonal entries
    diag = tf.linalg.diag(tf.ones(batch_size, dtype=tf.bool))
    diag_mask = t.cast(tf.logical_not(diag), dtype=query_embeddings.dtype)
    negexpd = tf.math.multiply(negexpd, diag_mask)

    # creating mask to sample same class neighboorhood
    pos_mask, _ = build_masks(
        query_labels,
        key_labels,
        batch_size=batch_size,
        remove_diagonal=remove_diagonal,
    )
    pos_mask = t.cast(pos_mask, dtype=query_embeddings.dtype)

    # all class neighborhood
    alcn = tf.reduce_sum(negexpd, axis=1)

    # same class neighborhood
    sacn = tf.reduce_sum(tf.math.multiply(negexpd, pos_mask), axis=1)

    # exclude examples with unique class from loss calculation
    excl = tf.math.not_equal(tf.reduce_sum(pos_mask, axis=1), tf.zeros(batch_size))
    excl = t.cast(excl, dtype=query_embeddings.dtype)

    loss = tf.math.divide(sacn, alcn)
    loss = -tf.multiply(tf.math.log(eps + loss), excl)

    return loss
//  [python]
def soft_nn_loss(
    query_labels: IntTensor,
    query_embeddings: FloatTensor,
    key_labels: IntTensor,
    key_embeddings: FloatTensor,
    distance: Distance,
    temperature: float,
    remove_diagonal: bool = True,
) -> Any:
    """Computes the soft nearest neighbors loss.

    Args:
        query_labels: labels associated with the query embed.
        query_embeddings: Embedded query examples.
        key_labels: labels associated with the key embed.
        key_embeddings: Embedded key examples.
        distance: Which distance function to use to compute the pairwise.
        temperature: Controls relative importance given to the pair of points.
        remove_diagonal: Bool. If True, will set diagonal to False in positive pair mask

    Returns:
        loss: loss value for the current batch.
    """

    batch_size = tf.size(query_labels)
    eps = tf.cast(1e-9, dtype=query_embeddings.dtype)

    pairwise_dist = distance(query_embeddings, key_embeddings)
    pairwise_dist = pairwise_dist / temperature
    negexpd = tf.math.exp(-pairwise_dist)

    # Mask out diagonal entries
    diag = tf.linalg.diag(tf.ones(batch_size, dtype=tf.bool))
    diag_mask = tf.cast(tf.logical_not(diag), dtype=query_embeddings.dtype)
    negexpd = tf.math.multiply(negexpd, diag_mask)

    # creating mask to sample same class neighboorhood
    pos_mask, _ = build_masks(
        query_labels,
        key_labels,
        batch_size=batch_size,
        remove_diagonal=remove_diagonal,
    )
    pos_mask = tf.cast(pos_mask, dtype=query_embeddings.dtype)

    # all class neighborhood
    alcn = tf.reduce_sum(negexpd, axis=1)

    # same class neighborhood
    sacn = tf.reduce_sum(tf.math.multiply(negexpd, pos_mask), axis=1)

    # exclude examples with unique class from loss calculation
    excl = tf.math.not_equal(tf.reduce_sum(pos_mask, axis=1), tf.zeros(batch_size))
    excl = tf.cast(excl, dtype=query_embeddings.dtype)

    loss = tf.math.divide(sacn, alcn)
    loss = -tf.multiply(tf.math.log(eps + loss), excl)

    return loss
//  [python]
def _add_positional_embedding(self, inputs):
    """Adds positional embeddings to the inputs tensor."""
    # Compute the locations using meshgrid.
    b, t, h, w = _get_shape(inputs)[:4]

    mesh = tf.meshgrid(tf.range(t), tf.range(h), tf.range(w), indexing='ij')
    position = tf.tile(
            tf.expand_dims(tf.stack(mesh, axis=-1), axis=0), [b, 1, 1, 1, 1])

    # Make the positions relative to center point
    # The mean of all position coordinates would be the center point anyway.
    center_position = tf.reduce_mean(position, axis=[1, 2, 3], keepdims=True)
    position -= center_position

    # Apply learneable layers.
    temporal_position = position[..., :1]
    for mlp in self._temporal_mlp:
      temporal_position = mlp(temporal_position)
    spatial_position = position[..., 1:]
    for mlp in self._spatial_mlp:
      spatial_position = mlp(spatial_position)

    return tf.concat([inputs, temporal_position, spatial_position], axis=-1)
//  [python]
def _add_positional_embedding(self, inputs):
    """Adds positional embeddings to the inputs tensor."""
    # Compute the locations using meshgrid.
    b, t, h, w = _get_shape(inputs)[:4]

    mesh = tf.meshgrid(tf.range(t), tf.range(h), tf.range(w), indexing='ij')
    position = tf.cast(
        tf.tile(
            tf.expand_dims(tf.stack(mesh, axis=-1), axis=0), [b, 1, 1, 1, 1]),
        tf.float32)

    # Make the positions relative to center point
    # The mean of all position coordinates would be the center point anyway.
    center_position = tf.reduce_mean(position, axis=[1, 2, 3], keepdims=True)
    position -= center_position

    # Apply learneable layers.
    temporal_position = position[..., :1]
    for mlp in self._temporal_mlp:
      temporal_position = mlp(temporal_position)
    spatial_position = position[..., 1:]
    for mlp in self._spatial_mlp:
      spatial_position = mlp(spatial_position)

    return tf.concat([inputs, temporal_position, spatial_position], axis=-1)
//  [python]
def call(self, inputs: tf.Tensor):
    """Calls text.BertTokenizer on inputs.

    Args:
      inputs: A string Tensor of shape [batch_size].

    Returns:
      One or three of RaggedTensors if tokenize_with_offsets is False or True,
      respectively. These are
      tokens: A RaggedTensor of shape [batch_size, (words), (pieces_per_word)]
        and type int32. tokens[i,j,k] contains the k-th wordpiece of the
        j-th word in the i-th input.
      start_offsets, limit_offsets: If tokenize_with_offsets is True,
        RaggedTensors of type int64 with the same indices as tokens.
        Element [i,j,k] contains the byte offset at the start, or past the
        end, resp., for the k-th wordpiece of the j-th word in the i-th input.
    """
    # Prepare to reshape the result to work around broken shape inference.
    batch_size = tf.shape(inputs)[0]
    def _reshape(rt):
      values = rt.values
      row_splits = rt.row_splits
      row_splits = tf.reshape(row_splits, [batch_size + 1])
      return tf.RaggedTensor.from_row_splits(values, row_splits)

    # Call the tokenizer.
    if self.tokenize_with_offsets:
      tokens, start_offsets, limit_offsets = (
          self._bert_tokenizer.tokenize_with_offsets(inputs))
      tokens = t.cast(tokens, dtype=tf.int32)
      return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
    else:
      tokens = self._bert_tokenizer.tokenize(inputs)
      tokens = t.cast(tokens, dtype=tf.int32)
      return _reshape(tokens)
//  [python]
def call(self, inputs: tf.Tensor):
    """Calls text.BertTokenizer on inputs.

    Args:
      inputs: A string Tensor of shape [batch_size].

    Returns:
      One or three of RaggedTensors if tokenize_with_offsets is False or True,
      respectively. These are
      tokens: A RaggedTensor of shape [batch_size, (words), (pieces_per_word)]
        and type int32. tokens[i,j,k] contains the k-th wordpiece of the
        j-th word in the i-th input.
      start_offsets, limit_offsets: If tokenize_with_offsets is True,
        RaggedTensors of type int64 with the same indices as tokens.
        Element [i,j,k] contains the byte offset at the start, or past the
        end, resp., for the k-th wordpiece of the j-th word in the i-th input.
    """
    # Prepare to reshape the result to work around broken shape inference.
    batch_size = tf.shape(inputs)[0]
    def _reshape(rt):
      values = rt.values
      row_splits = rt.row_splits
      row_splits = tf.reshape(row_splits, [batch_size + 1])
      return tf.RaggedTensor.from_row_splits(values, row_splits)

    # Call the tokenizer.
    if self.tokenize_with_offsets:
      tokens, start_offsets, limit_offsets = (
          self._bert_tokenizer.tokenize_with_offsets(inputs))
      tokens = tf.cast(tokens, dtype=tf.int32)
      return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
    else:
      tokens = self._bert_tokenizer.tokenize(inputs)
      tokens = tf.cast(tokens, dtype=tf.int32)
      return _reshape(tokens)
//  [python]
def multiword_pool_qkv_attention_global(inputs, embed, segment_ids, segment_locations, n_heads, keep_prob, size, is_training, pool_fn,
                                        combine="concatenate",
                                        attend_over_embed=False, log_tensorboard=False):
    # obtain Time(unique) x Batch x Hidden
    pooled = pool_fn(inputs, segment_ids, segment_locations, is_training=is_training, embed=embed)
    # now also do attention over the original inputs, but only from the point of
    # view of the pooled locations
    batch_major_inputs = tf.transpose(tf.concat([inputs, embed], axis=-1) if attend_over_embed else inputs, [1, 0, 2])
    batch_major_pooled = tf.transpose(pooled, [1, 0, 2])
    q, k, v = compute_qkv(batch_major_pooled, batch_major_inputs, size, size, size, size,
                          q_padding="VALID", kv_padding="VALID",
                          weight_noise=0.0, is_training=is_training)
    dropout_rate = tf.cond(is_training, lambda: 1.0 - keep_prob, lambda: 0.0)
    seqlen = tf.get_collection(SEQUENCE_LENGTHS)[0]

    # Convert # Batch x K-Time
    # into
    # Batch x Heads x Q-time x K-Time
    bias = -50.0 * tf.expand_dims(tf.expand_dims(tf.logical_not(tf.sequence_mask(seqlen)), 1), 1)
    x = multihead_attention_internal(q, k, v, dropout_rate, bias,
        num_heads=n_heads, total_key_depth=size, image_shapes=None)
    weights = multihead_attention_internal_weights(q, k, dropout_rate, bias,
        num_heads=n_heads, total_key_depth=size, image_shapes=None)
    tf.add_to_collection(ATTENTION_WEIGHTS, weights)
    if log_tensorboard:
        for head_idx in range(n_heads):
            tf_logger.test_image_summary_bw("multihead_attention_head{}".format(head_idx), weights[0, head_idx, :, :, None])
    q_global = tf.get_variable(name="GlobalQuery", shape=[1, 1, size])
    x_global = multihead_attention_internal(q_global, k, v, dropout_rate, bias,
         num_heads=n_heads, total_key_depth=size, image_shapes=None)
    weights = multihead_attention_internal_weights(q_global, k, dropout_rate, bias,
        num_heads=n_heads, total_key_depth=size, image_shapes=None)
    tf.add_to_collection(ATTENTION_WEIGHTS, weights)
    x = tf.transpose(x, [1, 0, 2])
    x_global = tf.tile(tf.transpose(x_global, [1, 0, 2]), [tf.shape(pooled)[0], 1, 1])
    if combine == "concatenate":
        return tf.concat([pooled, x, x_global], axis=-1)
    elif combine == "replace":
        return tf.concat([x, x_global], axis=-1)
    else:
        raise ValueError()
//  [python]
def multiword_pool_qkv_attention_global(inputs, embed, segment_ids, segment_locations, n_heads, keep_prob, size, is_training, pool_fn,
                                        combine="concatenate",
                                        attend_over_embed=False, log_tensorboard=False):
    # obtain Time(unique) x Batch x Hidden
    pooled = pool_fn(inputs, segment_ids, segment_locations, is_training=is_training, embed=embed)
    # now also do attention over the original inputs, but only from the point of
    # view of the pooled locations
    batch_major_inputs = tf.transpose(tf.concat([inputs, embed], axis=-1) if attend_over_embed else inputs, [1, 0, 2])
    batch_major_pooled = tf.transpose(pooled, [1, 0, 2])
    q, k, v = compute_qkv(batch_major_pooled, batch_major_inputs, size, size, size, size,
                          q_padding="VALID", kv_padding="VALID",
                          weight_noise=0.0, is_training=is_training)
    dropout_rate = tf.cond(is_training, lambda: 1.0 - keep_prob, lambda: 0.0)
    seqlen = tf.get_collection(SEQUENCE_LENGTHS)[0]

    # Convert # Batch x K-Time
    # into
    # Batch x Heads x Q-time x K-Time
    bias = -50.0 * tf.cast(tf.expand_dims(tf.expand_dims(tf.logical_not(tf.sequence_mask(seqlen)), 1), 1), tf.float32)
    x = multihead_attention_internal(q, k, v, dropout_rate, bias,
        num_heads=n_heads, total_key_depth=size, image_shapes=None)
    weights = multihead_attention_internal_weights(q, k, dropout_rate, bias,
        num_heads=n_heads, total_key_depth=size, image_shapes=None)
    tf.add_to_collection(ATTENTION_WEIGHTS, weights)
    if log_tensorboard:
        for head_idx in range(n_heads):
            tf_logger.test_image_summary_bw("multihead_attention_head{}".format(head_idx), weights[0, head_idx, :, :, None])
    q_global = tf.get_variable(name="GlobalQuery", shape=[1, 1, size])
    x_global = multihead_attention_internal(q_global, k, v, dropout_rate, bias,
         num_heads=n_heads, total_key_depth=size, image_shapes=None)
    weights = multihead_attention_internal_weights(q_global, k, dropout_rate, bias,
        num_heads=n_heads, total_key_depth=size, image_shapes=None)
    tf.add_to_collection(ATTENTION_WEIGHTS, weights)
    x = tf.transpose(x, [1, 0, 2])
    x_global = tf.tile(tf.transpose(x_global, [1, 0, 2]), [tf.shape(pooled)[0], 1, 1])
    if combine == "concatenate":
        return tf.concat([pooled, x, x_global], axis=-1)
    elif combine == "replace":
        return tf.concat([x, x_global], axis=-1)
    else:
        raise ValueError()
//  [python]
def validation_step(self, inputs, model, metrics=None):
    if self.task_config.model.supervised_head is None:
      raise ValueError(
          'Skipping eval during pretraining without supervised head.')

    features, labels = inputs
    if self.task_config.evaluation.one_hot:
      num_classes = self.task_config.model.supervised_head.num_classes
      labels = tf.one_hot(labels, num_classes)

    outputs = model(
        features, training=False)[simclr_model.SUPERVISED_OUTPUT_KEY]
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

    logs = {self.loss: 0}

    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})

    return logs
//  [python]
pred = RNN(x, weights, biases)
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
train_op = tf.train.AdamOptimizer(lr).minimize(cost)

correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.to_float(correct_pred))
//  [python]
pred = RNN(x, weights, biases)
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
train_op = tf.train.AdamOptimizer(lr).minimize(cost)

correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    features, labels = inputs

    # To do a sanity check that we absolutely use no labels when pretraining, we
    # can set the labels here to zero.
    if self.task_config.train_data.input_set_label_to_zero:
      labels *= 0

    if (self.task_config.model.supervised_head is not None and
        self.task_config.evaluation.one_hot):
      num_classes = self.task_config.model.supervised_head.num_classes
      labels = tf.one_hot(labels, num_classes)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)

      scaled_loss = losses['total_loss'] / num_replicas
      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    logging.info('Trainable variables:')
    for var in tvars:
      logging.info(var.name)
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    for m in metrics:
      m.update_state(losses[m.name])
      logs.update({m.name: m.result()})

    return logs
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    features, labels = inputs

    # To do a sanity check that we absolutely use no labels when pretraining, we
    # can set the labels here to zero.
    if self.task_config.train_data.input_set_label_to_zero:
      labels *= 0

    if (self.task_config.model.supervised_head is not None and
        self.task_config.evaluation.one_hot):
      num_classes = self.task_config.model.supervised_head.num_classes
      labels = tf.one_hot(labels, num_classes)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)

      scaled_loss = losses['total_loss'] / num_replicas
      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    logging.info('Trainable variables:')
    for var in tvars:
      logging.info(var.name)
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    for m in metrics:
      m.update_state(losses[m.name])
      logs.update({m.name: m.result()})

    return logs
//  [python]
def build_losses(self,
                   labels,
                   model_outputs,
                   aux_losses=None) -> Dict[str, tf.Tensor]:
    # Compute contrastive relative loss
    con_losses_obj = contrastive_losses.ContrastiveLoss(
        projection_norm=self.task_config.loss.projection_norm,
        temperature=self.task_config.loss.temperature)
    # The projection outputs from model has the size of
    # (2 * bsz, project_dim)
    projection_outputs = model_outputs[simclr_model.PROJECTION_OUTPUT_KEY]
    projection1, projection2 = tf.split(projection_outputs, 2, 0)
    contrast_loss, (contrast_logits, contrast_labels) = con_losses_obj(
        projection1=projection1, projection2=projection2)

    contrast_accuracy = tf.equal(
        tf.argmax(contrast_labels, axis=1), tf.argmax(contrast_logits, axis=1))
    contrast_accuracy = tf.reduce_mean(tf.to_float(contrast_accuracy))

    contrast_prob = tf.nn.softmax(contrast_logits)
    contrast_entropy = -tf.reduce_mean(
        tf.reduce_sum(contrast_prob * tf.math.log(contrast_prob + 1e-8), -1))

    model_loss = contrast_loss

    losses = {
        'contrast_loss': contrast_loss,
        'contrast_accuracy': contrast_accuracy,
        'contrast_entropy': contrast_entropy
    }

    if self.task_config.model.supervised_head is not None:
      outputs = model_outputs[simclr_model.SUPERVISED_OUTPUT_KEY]
      labels = tf.concat([labels, labels], 0)

      if self.task_config.evaluation.one_hot:
        sup_loss = tf_keras.losses.CategoricalCrossentropy(
            from_logits=True, reduction=tf_keras.losses.Reduction.NONE)(labels,
                                                                        outputs)
      else:
        sup_loss = tf_keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf_keras.losses.Reduction.NONE)(labels,
                                                                        outputs)
      sup_loss = tf.reduce_mean(sup_loss)

      label_acc = tf.equal(
          tf.argmax(labels, axis=1), tf.argmax(outputs, axis=1))
      label_acc = tf.reduce_mean(tf.to_float(label_acc))

      model_loss = contrast_loss + sup_loss

      losses.update({
          'accuracy': label_acc,
          'supervised_loss': sup_loss,
      })

    total_loss = model_loss
    if aux_losses:
      reg_loss = tf.reduce_sum(aux_losses)
      total_loss = model_loss + reg_loss

    losses['total_loss'] = total_loss

    return losses
//  [python]
def build_losses(self,
                   labels,
                   model_outputs,
                   aux_losses=None) -> Dict[str, tf.Tensor]:
    # Compute contrastive relative loss
    con_losses_obj = contrastive_losses.ContrastiveLoss(
        projection_norm=self.task_config.loss.projection_norm,
        temperature=self.task_config.loss.temperature)
    # The projection outputs from model has the size of
    # (2 * bsz, project_dim)
    projection_outputs = model_outputs[simclr_model.PROJECTION_OUTPUT_KEY]
    projection1, projection2 = tf.split(projection_outputs, 2, 0)
    contrast_loss, (contrast_logits, contrast_labels) = con_losses_obj(
        projection1=projection1, projection2=projection2)

    contrast_accuracy = tf.equal(
        tf.argmax(contrast_labels, axis=1), tf.argmax(contrast_logits, axis=1))
    contrast_accuracy = tf.reduce_mean(tf.cast(contrast_accuracy, tf.float32))

    contrast_prob = tf.nn.softmax(contrast_logits)
    contrast_entropy = -tf.reduce_mean(
        tf.reduce_sum(contrast_prob * tf.math.log(contrast_prob + 1e-8), -1))

    model_loss = contrast_loss

    losses = {
        'contrast_loss': contrast_loss,
        'contrast_accuracy': contrast_accuracy,
        'contrast_entropy': contrast_entropy
    }

    if self.task_config.model.supervised_head is not None:
      outputs = model_outputs[simclr_model.SUPERVISED_OUTPUT_KEY]
      labels = tf.concat([labels, labels], 0)

      if self.task_config.evaluation.one_hot:
        sup_loss = tf_keras.losses.CategoricalCrossentropy(
            from_logits=True, reduction=tf_keras.losses.Reduction.NONE)(labels,
                                                                        outputs)
      else:
        sup_loss = tf_keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf_keras.losses.Reduction.NONE)(labels,
                                                                        outputs)
      sup_loss = tf.reduce_mean(sup_loss)

      label_acc = tf.equal(
          tf.argmax(labels, axis=1), tf.argmax(outputs, axis=1))
      label_acc = tf.reduce_mean(tf.cast(label_acc, tf.float32))

      model_loss = contrast_loss + sup_loss

      losses.update({
          'accuracy': label_acc,
          'supervised_loss': sup_loss,
      })

    total_loss = model_loss
    if aux_losses:
      reg_loss = tf.reduce_sum(aux_losses)
      total_loss = model_loss + reg_loss

    losses['total_loss'] = total_loss

    return losses
//  [python]
def positional_encoding(position, embed_dim):
  angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(embed_dim)[np.newaxis, :],
                          embed_dim)

  # apply sin to even indices in the array; 2i
  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

  # apply cos to odd indices in the array; 2i+1
  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

  pos_encoding = angle_rads[np.newaxis, ...]

  return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
    dk = tf.to_float(tf.shape(k)[-1])
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e9)  
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
    output = tf.matmul(attention_weights, v)
    return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e9)  
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
    output = tf.matmul(attention_weights, v)
    return output, attention_weights
//  [python]
def validation_step(self, inputs, model, metrics=None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    if self.task_config.loss.one_hot:
      num_classes = self.task_config.model.supervised_head.num_classes
      labels = tf.one_hot(labels, num_classes)

    outputs = self.inference_step(features,
                                  model)[simclr_model.SUPERVISED_OUTPUT_KEY]
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    loss = self.build_losses(
        model_outputs=outputs, labels=labels, aux_losses=model.losses)

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def call(self, x, training, mask=None):

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, embed_dim)
        x *= tf.sqrt(tf.to_float(self.embed_dim))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x
//  [python]
def call(self, x, training, mask=None):

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, embed_dim)
        x *= tf.math.sqrt(tf.cast(self.embed_dim, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    if self.task_config.loss.one_hot:
      num_classes = self.task_config.model.supervised_head.num_classes
      labels = tf.one_hot(labels, num_classes)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(
          features, training=True)[simclr_model.SUPERVISED_OUTPUT_KEY]
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      loss = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    logging.info('Trainable variables:')
    for var in tvars:
      logging.info(var.name)
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def variable_summaries(var,name):
    if var.dtype != tf.float32:
        var = tf.cast(var, tf.float32)
    mean = tf.reduce_mean(var)
    tf.summary.scalar(name+'/max', tf.reduce_max(var))
    tf.summary.scalar(name+'/min', tf.reduce_min(var))
    tf.summary.scalar(name+'/mean', mean)
    stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
    tf.summary.scalar(name+'/stddev', stddev)
    tf.summary.histogram(name+'/hisogram', var)
//  [python]
def detection_image_summary_with_croped_mask(images,
                           boxes,
                           *args,
                           instance_masks=None,
                           **kwargs):
    if instance_masks is not None:
        shape = btf.combined_static_and_dynamic_shape(images)
        assert len(instance_masks.get_shape()) == 4,"error mask dims"
        assert len(boxes.get_shape()) == 3,"error boxes dims"
        assert len(images.get_shape()) == 4,"error images dims"
        instance_masks = imv.batch_tf_get_fullsize_mask(boxes=boxes,
                                                        masks=instance_masks,
                                                        size=shape[1:3]
                                                        )
        if instance_masks.dtype != tf.uint8 and instance_masks.dtype != tf.bool:
            instance_masks = tf.to_float(tf.cast(instance_masks > 0.5, tf.float32))
    return detection_image_summary(images,boxes,*args,instance_masks=instance_masks,**kwargs)
//  [python]
def detection_image_summary_with_croped_mask(images,
                           boxes,
                           *args,
                           instance_masks=None,
                           **kwargs):
    if instance_masks is not None:
        shape = btf.combined_static_and_dynamic_shape(images)
        assert len(instance_masks.get_shape()) == 4,"error mask dims"
        assert len(boxes.get_shape()) == 3,"error boxes dims"
        assert len(images.get_shape()) == 4,"error images dims"
        instance_masks = imv.batch_tf_get_fullsize_mask(boxes=boxes,
                                                        masks=instance_masks,
                                                        size=shape[1:3]
                                                        )
        if instance_masks.dtype != tf.uint8 and instance_masks.dtype != tf.bool:
            instance_masks = tf.cast(instance_masks > 0.5, tf.float32)
    return detection_image_summary(images,boxes,*args,instance_masks=instance_masks,**kwargs)
//  [python]
def _encode(
        self, rationals: Union[tf.Tensor, np.ndarray], apply_scaling: bool
    ) -> Union[tf.Tensor, np.ndarray]:
        """
        Encode tensor of rational numbers into tensor of ring elements. Output is
        of same type as input to allow function to be used for constants.
        """

        with tf.name_scope("encode"):

            if isinstance(rationals, np.ndarray):
                if apply_scaling:
                    scaled = rationals * self.fixedpoint_config.scaling_factor
                    scaled = np.array(scaled).astype(np.float64)
                else:
                    scaled = rationals

            elif isinstance(rationals, tf.Tensor):
                if apply_scaling:
                    scaled = t.cast(rationals, tf.float64)
                    scaled = scaled * self.fixedpoint_config.scaling_factor
                else:
                    scaled = rationals

            else:
                # give it a last try
                try:
                    if apply_scaling:
                        scaled = rationals * self.fixedpoint_config.scaling_factor
                        scaled = np.array(scaled, dtype=object)
                    else:
                        scaled = np.array(rationals, dtype=object)
                except:
                    raise TypeError(
                        "Don't know how to encode {}".format(type(rationals))
                    )

            return scaled
//  [python]
def _encode(
        self, rationals: Union[tf.Tensor, np.ndarray], apply_scaling: bool
    ) -> Union[tf.Tensor, np.ndarray]:
        """
        Encode tensor of rational numbers into tensor of ring elements. Output is
        of same type as input to allow function to be used for constants.
        """

        with tf.name_scope("encode"):

            if isinstance(rationals, np.ndarray):
                if apply_scaling:
                    scaled = rationals * self.fixedpoint_config.scaling_factor
                    scaled = np.array(scaled).astype(np.float64)
                else:
                    scaled = rationals

            elif isinstance(rationals, tf.Tensor):
                if apply_scaling:
                    scaled = tf.cast(rationals, tf.float64)
                    scaled = scaled * self.fixedpoint_config.scaling_factor
                else:
                    scaled = rationals

            else:
                # give it a last try
                try:
                    if apply_scaling:
                        scaled = rationals * self.fixedpoint_config.scaling_factor
                        scaled = np.array(scaled, dtype=object)
                    else:
                        scaled = np.array(rationals, dtype=object)
                except:
                    raise TypeError(
                        "Don't know how to encode {}".format(type(rationals))
                    )

            return scaled
//  [python]
def variable_summaries_v2(var, name,family=None):
    if var is None or name is None:
        return
    if var.dtype != tf.float32:
        var = tf.cast(var, tf.float32)
    if isSingleValueTensor(var):
        var = tf.reshape(var,())
        tf.summary.scalar(name, var,family=family)
        return
    mean = tf.reduce_mean(var)
    tf.summary.scalar(name+'/max', tf.reduce_max(var),family=family)
    tf.summary.scalar(name+'/min', tf.reduce_min(var),family=family)
    tf.summary.scalar(name+'/mean', mean,family=family)
    stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
    tf.summary.scalar(name+'/stddev', stddev,family=family)
    tf.summary.histogram(name+'/hisogram', var,family=family)
//  [python]
def feature_map_summary(feature_map,name="feature_map",max_outputs=None,min_max=None):
    data = feature_map[0]
    if data.dtype != tf.float32:
        data = tf.to_float(data)
    if max_outputs is None:
        max_outputs = feature_map.get_shape().as_list()[-1]
    data = tf.transpose(data,[2,0,1])
    data = tf.expand_dims(data,axis=-1)
    if min_max is None:
        min = tf.reduce_min(data)
        max = tf.reduce_max(data)
    else:
        min,max = min_max
    data = (data-min)/(max-min+1e-8)
    tf.summary.image(name,data,max_outputs=max_outputs)
//  [python]
def feature_map_summary(feature_map,name="feature_map",max_outputs=None,min_max=None):
    data = feature_map[0]
    if data.dtype != tf.float32:
        data = tf.cast(data,tf.float32)
    if max_outputs is None:
        max_outputs = feature_map.get_shape().as_list()[-1]
    data = tf.transpose(data,[2,0,1])
    data = tf.expand_dims(data,axis=-1)
    if min_max is None:
        min = tf.reduce_min(data)
        max = tf.reduce_max(data)
    else:
        min,max = min_max
    data = (data-min)/(max-min+1e-8)
    tf.summary.image(name,data,max_outputs=max_outputs)
//  [python]
def semantic_summary_img(img,masks,colors=None,name="semantic"):
    ''''
    image: [height, width, 3] value range[0, 255]
    mask: [height, width,N] value range[0, 1]
    '''
    with tf.device("/cpu:0"):
        masks = tf.transpose(masks,[2,0,1])
        if img.dtype != tf.uint8:
            min = tf.reduce_min(img)
            max = tf.reduce_max(img)
            img = (img-min)*255/(max-min+1e-8)
            img = t.cast(img,tf.uint8)
        if colors is None:
            colors = smv.STANDARD_COLORS
        colors = tf.convert_to_tensor(colors)
        img = smv.tf_draw_masks_on_image(img,masks,color=colors)
    return img
//  [python]
def semantic_summary_img(img,masks,colors=None,name="semantic"):
    ''''
    image: [height, width, 3] value range[0, 255]
    mask: [height, width,N] value range[0, 1]
    '''
    with tf.device("/cpu:0"):
        masks = tf.transpose(masks,[2,0,1])
        if img.dtype != tf.uint8:
            min = tf.reduce_min(img)
            max = tf.reduce_max(img)
            img = (img-min)*255/(max-min+1e-8)
            img = tf.cast(img,tf.uint8)
        if colors is None:
            colors = smv.STANDARD_COLORS
        colors = tf.convert_to_tensor(colors)
        img = smv.tf_draw_masks_on_image(img,masks,color=colors)
    return img
//  [python]
def schedule_at_step(step,
                     vals,
                     boundaries = (),
                     warmup_steps = 0,
                     ):
  """Computes the schedule value at a given step `step`.

  Args:
    step: The step to compute the schedule value at.
    vals: Sequence of values.
    boundaries: Locations where the schedule changes between values in `vals`.
      If empty, `vals` should be a sequence with exactly one element.
    warmup_steps:  Number of steps at the beginning of training to use as
      warmup. Set to non-positive to disable.

  Returns:
    The computed schedule value at step `step`.
  """
  step = tf.convert_to_tensor(step)
  value = piecewise_constant_schedule(step, boundaries, vals)
  if warmup_steps > 0:
    # Applies linear warmup, over the first `warmup_steps` steps.
    value *= tf.minimum(1., (tf.cast(step, tf.float32) + 1) / warmup_steps)
  return value
//  [python]
def summary_image_with_box(image,
                                bboxes,
                                classes,
                                scores=None,
                                category_index=None,
                                max_boxes_to_draw=100,
                                name='summary_image_with_box',scale=True):
    with tf.name_scope(name):
        if scale:
            if (image.dtype==tf.float32) or (image.dtype==tf.float16) or (image.dtype==tf.float64):
                #floatpoint data value range is [-1,1]
                image = (image+1.0)*127.5
                image = tf.clip_by_value(image,clip_value_min=0.,clip_value_max=255.)
                image = tf.cast(image,dtype=tf.uint8)
        if image.get_shape().ndims == 3:
            image = tf.expand_dims(image,axis=0)
        if image.dtype is not tf.uint8:
            image = tf.cast(image,tf.uint8)
        if bboxes.get_shape().ndims == 2:
            bboxes = tf.expand_dims(bboxes, 0)
        if scores is None:
            scores = tf.ones(shape=tf.shape(bboxes)[:2],dtype=tf.float32)
        elif scores.get_shape().ndims ==1:
            scores = tf.expand_dims(scores,axis=0)
        if category_index is None:
            category_index = {}
            for i in range(100):
                category_index[i] = {"name":f"{i}"}
        if classes.get_shape().ndims == 1:
            classes = tf.expand_dims(classes,axis=0)
        image_with_box = smv.draw_bounding_boxes_on_image_tensors(image,bboxes,classes,scores,
                                                              category_index=category_index,
                                                              max_boxes_to_draw=max_boxes_to_draw)
        tf.summary.image(name, image_with_box)
//  [python]
def decode_image(img_bytes):
    pixels = tf.image.decode_image(img_bytes, channels=3)
    return tf.cast(pixels, tf.uint8)
//  [python]
def _parse_example_proto(self, example_serialized):
        """Parses an Example proto containing a training example of an image.
        The output of the build_image_data.py image preprocessing script is a
        dataset containing serialized Example protocol buffers. Each Example
        proto contains the following fields:
        image/height: 462
        image/width: 581
        image/colorspace: 'RGB'
        image/channels: 3
        image/class/label: 615
        image/class/synset: 'n03623198'
        image/class/text: 'knee pad'
        image/format: 'JPEG'
        image/filename: 'ILSVRC2012_val_00041207.JPEG'
        image/encoded: <JPEG encoded string>

        Args:
          example_serialized (tf.string): Scalar Tensor tf.string containing a
          serialized Example protocol buffer.

        Returns:
          tupel: Tupel of image_buffer (tf.string) containing the contents of a
          JPEG file, the label (tf.int32) containing the label and text
          (tf.string) containing the human-readable label.
        """
        # Dense features in Example proto.
        feature_map = {
            'image/encoded':
            tf.FixedLenFeature([], dtype=tf.string, default_value=''),
            'image/class/label':
            tf.FixedLenFeature([1], dtype=tf.int64, default_value=-1),
            'image/class/text':
            tf.FixedLenFeature([], dtype=tf.string, default_value=''),
        }

        features = tf.parse_single_example(example_serialized, feature_map)
        label = tf.to_int32(features['image/class/label'])

        return features['image/encoded'], label, features['image/class/text']
//  [python]
def _parse_example_proto(self, example_serialized):
        """Parses an Example proto containing a training example of an image.
        The output of the build_image_data.py image preprocessing script is a
        dataset containing serialized Example protocol buffers. Each Example
        proto contains the following fields:
        image/height: 462
        image/width: 581
        image/colorspace: 'RGB'
        image/channels: 3
        image/class/label: 615
        image/class/synset: 'n03623198'
        image/class/text: 'knee pad'
        image/format: 'JPEG'
        image/filename: 'ILSVRC2012_val_00041207.JPEG'
        image/encoded: <JPEG encoded string>

        Args:
          example_serialized (tf.string): Scalar Tensor tf.string containing a
          serialized Example protocol buffer.

        Returns:
          tupel: Tupel of image_buffer (tf.string) containing the contents of a
          JPEG file, the label (tf.int32) containing the label and text
          (tf.string) containing the human-readable label.
        """
        # Dense features in Example proto.
        feature_map = {
            'image/encoded':
            tf.FixedLenFeature([], dtype=tf.string, default_value=''),
            'image/class/label':
            tf.FixedLenFeature([1], dtype=tf.int64, default_value=-1),
            'image/class/text':
            tf.FixedLenFeature([], dtype=tf.string, default_value=''),
        }

        features = tf.parse_single_example(example_serialized, feature_map)
        label = tf.cast(features['image/class/label'], dtype=tf.int32)

        return features['image/encoded'], label, features['image/class/text']
//  [python]
def process_reals(x, lod, mirror_augment, drange_data, drange_net):
    with tf.name_scope('ProcessReals'):
        with tf.name_scope('DynamicRange'):
            x = tf.cast(x, tf.float32)
            x = misc.adjust_dynamic_range(x, drange_data, drange_net)
        if mirror_augment:
            with tf.name_scope('MirrorAugment'):
                s = tf.shape(x)
                mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)
                mask = tf.tile(mask, [1, s[1], s[2], s[3]])
                x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))
        with tf.name_scope('FadeLOD'): # Smooth crossfade between consecutive levels-of-detail.
            s = tf.shape(x)
            y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])
            y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
            y = tf.tile(y, [1, 1, 1, 2, 1, 2])
            y = tf.reshape(y, [-1, s[1], s[2], s[3]])
            x = tflib.lerp(x, y, lod - tf.floor(lod))
        with tf.name_scope('UpscaleLOD'): # Upscale to match the expected input/output size of the networks.
            s = tf.shape(x)
            factor = tf.cast(2 ** tf.floor(lod), tf.int32)
            x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
            x = tf.tile(x, [1, 1, 1, factor, 1, factor])
            x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
        return x