def coreference_resolution_loss(
    mention_encodings: Array,
    mention_batch_positions: Array,
    mention_target_ids: Array,
    batch_size: int,
    mode: str,
    mention_target_is_masked: Optional[Array] = None,
    metrics_prefix: str = '') -> Tuple[float, Dict[str, Dict[str, float]]]:
  """Compute coreference resolution loss over global batch (across all devices).

  Args:
    mention_encodings: [n_mentions, hidden_dim] mention encodings to be used for
      computing the loss.
    mention_batch_positions: [n_mentions] position of a mention in its batch.
    mention_target_ids: [n_mentions] IDs of mentions.
    batch_size: batch size per device.
    mode: how to compute the scores -- using dot product ('dot'), dot product
      divided by the sqrt root of the hidden dim ('dot_sqrt') or cosine
      similarity ('cos').
    mention_target_is_masked: [n_mentions] whether a mention has been masked. If
      it is provided the function computes additional metrics (accuracy) for
      masked and non masked mentions separately.
    metrics_prefix: optional prefix for all metric names.

  Returns:
    Loss and a dictionary of metrics dictionaries.
  """
  hidden_dim = mention_encodings.shape[-1]
  all_mention_encodings, all_mention_target_ids = get_all_encodings_and_target_ids(
      mention_encodings, mention_target_ids)

  (local_mention_batch_positions, global_mention_batch_positions
  ) = mention_utils.get_globally_consistent_batch_positions(
      mention_batch_positions, batch_size)

  positive_mask, negative_mask = build_coref_positive_negative_mask(
      local_mention_batch_positions, global_mention_batch_positions,
      mention_target_ids, all_mention_target_ids)

  # We compute a similarity between encodings of mentions from the
  # current device VS encodings of mentions gathered from all devices.
  unnormalized_scores = jnp.einsum('ld,gd->lg', mention_encodings,
                                   all_mention_encodings)

  # Set score to 0 if one of the mentions is PAD. This doesn't affect the loss
  # since we pass weights and positives / negatives masks explicitly.
  # However, it helps to compute various statistics like average scores
  # while ignoring scores for PAD mentions.
  mention_is_not_pad = jnp.expand_dims(mention_target_ids > 0, 1)
  mention_is_not_pad = mention_is_not_pad.astype(unnormalized_scores.dtype)
  all_mention_is_not_pad = jnp.expand_dims(all_mention_target_ids > 0, 0)
  all_mention_is_not_pad = all_mention_is_not_pad.astype(
      unnormalized_scores.dtype)
  unnormalized_scores = unnormalized_scores * mention_is_not_pad
  unnormalized_scores = unnormalized_scores * all_mention_is_not_pad

  # Norms of encodings of mentions from the current device.
  mention_encodings_norm = jnp.linalg.norm(mention_encodings, axis=-1)

  scores = unnormalized_scores
  if mode == 'dot':
    pass
  elif mode == 'dot_sqrt':
    scores /= jnp.sqrt(hidden_dim)
  elif mode == 'cos':
    # The cosine similarity is computed as dot product divided by norms of
    # both vectors.
    # Norms of encodings of mentions from all device.
    all_mention_encodings_norm = jnp.linalg.norm(all_mention_encodings, axis=-1)
    scores /= (_SMALL_NUMBER + jnp.expand_dims(mention_encodings_norm, 1))
    scores /= (_SMALL_NUMBER + jnp.expand_dims(all_mention_encodings_norm, 0))
  else:
    raise ValueError('Unknown coreference resolution mode: ' + mode)

  (loss, metrics, (acc_per_sample, weights_per_sample)
  ) = metric_utils.compute_cross_entropy_loss_with_positives_and_negatives_masks(
      scores, positive_mask, negative_mask, mention_target_ids != 0)

  num_positives = positive_mask.astype(jnp.float32).sum(1)
  num_positives = jnp.dot(num_positives, weights_per_sample)

  num_negatives = negative_mask.astype(jnp.float32).sum(1)
  num_negatives = jnp.dot(num_negatives, weights_per_sample)

  avg_score = scores.sum(axis=-1) / (
      all_mention_is_not_pad.sum() + _SMALL_NUMBER)
  avg_unnorm_score = unnormalized_scores.sum(axis=-1) / (
      all_mention_is_not_pad.sum() + _SMALL_NUMBER)

  metrics.update({
      'n_positive': num_positives,
      'n_negative': num_negatives,
      'avg_norm': jnp.dot(mention_encodings_norm, weights_per_sample),
      'avg_score': jnp.dot(avg_score, weights_per_sample),
      'avg_unnorm_score': jnp.dot(avg_unnorm_score, weights_per_sample),
  })

  final_metrics = {metrics_prefix + 'coref_resolution': metrics}
  if mention_target_is_masked is not None:
    final_metrics[metrics_prefix + 'coref_resolution_masked'] = {
        'acc':
            jnp.dot(acc_per_sample,
                    weights_per_sample * mention_target_is_masked),
        'denominator':
            jnp.dot(weights_per_sample, mention_target_is_masked),
    }
    final_metrics[metrics_prefix + 'coref_resolution_non_masked'] = {
        'acc':
            jnp.dot(acc_per_sample,
                    weights_per_sample * (1 - mention_target_is_masked)),
        'denominator':
            jnp.dot(weights_per_sample, (1 - mention_target_is_masked)),
    }

  return loss, final_metrics  # pytype: disable=bad-return-type  # jnp-type
Compute coreference resolution loss over global batch (across all devices).

  Args:
    mention_encodings: [n_mentions, hidden_dim] mention encodings to be used for
      computing the loss.
    mention_batch_positions: [n_mentions] position of a mention in its batch.
    mention_target_ids: [n_mentions] IDs of mentions.
    batch_size: batch size per device.
    mode: how to compute the scores -- using dot product ('dot'), dot product
      divided by the sqrt root of the hidden dim ('dot_sqrt') or cosine
      similarity ('cos').
    mention_target_is_masked: [n_mentions] whether a mention has been masked. If
      it is provided the function computes additional metrics (accuracy) for
      masked and non masked mentions separately.
    metrics_prefix: optional prefix for all metric names.

  Returns:
    Loss and a dictionary of metrics dictionaries.
  
language/mentionmemory/modules/mention_losses.py