//  [python]
def _compute_constrastive_loss(
      self,
      positive_lookup_index: tf.Tensor,
      inst_translated: tf.Tensor,
      inst_target: tf.Tensor,
      inst_mask: tf.Tensor,
      num_replicas: int = 1) -> Mapping[str, tf.Tensor]:
    """Computes constrastive loss.

    Args:
      positive_lookup_index: the index tensor to look-up the corresponding
        features in inst_target. In shape [B, N].
      inst_translated: a float tensor of shape [B, N, C] of translated instance
        features by the transformer head.
      inst_target: a float tensor of shape [B, N, C] of instance features on the
        target domain. Note that the order of inst_target is not necessarily
        matched to inst_translated.
      inst_mask: a boolean tensor of shape [B, N, 1] suggesting valid instances
        in inst_translated.
      num_replicas: the number of TPU replicas.

    Returns:
      loss_and_stats: a dictionary of loss and intermediate statistics.
    """
    b, n = inst_translated.shape.as_list()[:2]

    if num_replicas == 1:
      inst_target_large = inst_target
      b_large = tf.shape(inst_target_large)[0]
      labels_idx = tf.range(b)
    else:
      inst_target_large = tpu_cross_replica_concat(
          inst_target,
          num_replicas)
      b_large = tf.shape(inst_target_large)[0]
      # NOTE: make sure to use xla.replica_id() here and in
      # tpu_cross_replica_concat to consistently align the replica_id.
      # replicator.replica_id != xla.replica_id()
      replica_id = tf.cast(tf.cast(xla.replica_id(), tf.uint32), tf.int32)
      labels_idx = tf.range(b) + replica_id * b

    # [B, BL], 1 indicates positive batches.
    batch_masks = tf.one_hot(labels_idx, b_large)
    # [B, N, BL, N]
    batch_masks = tf.tile(batch_masks[:, None, :, None], [1, n, 1, n])

    # Construct negative examples.
    logits_negative = tf.einsum(
        'ijc,pqc->ijpq',
        inst_translated, inst_target_large) / self._temperature
    # Get negative statistics.
    negative_stats = self._get_negative_similarity_statistics(
        logits_negative, batch_masks, inst_mask)
    logits_negative = logits_negative - tf.cast(
        batch_masks, logits_negative.dtype) * _LARGE_NUM
    logits_negative = tf.reshape(logits_negative, [b * n, b_large * n])

    # Construct positive examples.
    inst_matched = tf.gather_nd(
        inst_target, positive_lookup_index, name='matched_inst')
    logits_positive = tf.einsum(
        'ijc,ijc->ij',
        inst_translated, inst_matched) / self._temperature
    # Get positive statistics.
    positive_stats = self._get_positive_similarity_statistics(
        logits_positive, inst_mask)
    logits_positive = tf.reshape(logits_positive, [b * n, 1])

    logits_all = tf.concat([logits_positive, logits_negative], axis=1)
    loss_pos = tf.reduce_logsumexp(logits_positive, 1)
    loss_all = tf.reduce_logsumexp(logits_all, 1)
    loss = (loss_all - loss_pos) * tf.reshape(inst_mask, [b * n])

    # Average across instances.
    loss = tf.math.divide_no_nan(
        tf.reduce_sum(loss), tf.reduce_sum(inst_mask))

    loss_and_stats = {'loss': loss}
    loss_and_stats.update(negative_stats)
    loss_and_stats.update(positive_stats)
    return loss_and_stats