def entity_linking_loss(mention_encodings: Array, entity_embeddings: Array,
                        mention_target_ids: Array,
                        mention_target_weights: Array, mode: str) -> Array:
  """Compute entity linking loss.

  Args:
    mention_encodings: [n_mentions, hidden_size] mention encodings to be used
      for computing the loss.
    entity_embeddings: [n_entities, hidden_size] entity embeddings table.
    mention_target_ids: [n_mentions] IDs of mentions.
    mention_target_weights: [n_mentions] per-mention weight for computing loss
      and metrics.
    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').

  Returns:
    Loss, a dictionary with metrics values, per sample infomation
    (a tuple of accuracy per mention and weight per mention).
  """
  scores = jnp.einsum('qd,ed->qe', mention_encodings, entity_embeddings)
  scores = scores.astype(jnp.float32)

  mention_encodings_norm = jnp.linalg.norm(mention_encodings, axis=-1)
  entity_embeddings_norm = jnp.linalg.norm(entity_embeddings, axis=-1)

  # The cosine similarity is computed as dot product divided by norms of
  # both vectors.
  cos_scores = scores
  cos_scores /= (_SMALL_NUMBER + jnp.expand_dims(mention_encodings_norm, 1))
  cos_scores /= (_SMALL_NUMBER + jnp.expand_dims(entity_embeddings_norm, 0))

  if mode == 'dot':
    pass
  elif mode == 'dot_sqrt':
    hidden_dim = mention_encodings.shape[1]
    scores /= jnp.sqrt(hidden_dim)
  elif mode == 'cos':
    scores = cos_scores
  else:
    raise ValueError('Unknown entity linking mode: ' + mode)

  mention_target_weights = mention_target_weights.astype(jnp.float32)

  loss, _ = metric_utils.compute_weighted_cross_entropy(
      scores, mention_target_ids, mention_target_weights, inputs_are_prob=False)

  acc_per_mention = jnp.equal(jnp.argmax(scores, axis=-1), mention_target_ids)

  acc_per_mention = acc_per_mention * mention_target_weights

  n_mentions = mention_target_ids.shape[0]
  cos_per_mention = cos_scores[jnp.arange(n_mentions), mention_target_ids]
  cos_per_mention = cos_per_mention * mention_target_weights

  metrics = {
      'loss': loss,
      'acc': acc_per_mention.sum(),
      'cos_sim': cos_per_mention.sum(),
      'denominator': mention_target_weights.sum()
  }
  return loss, metrics, (acc_per_mention, mention_target_weights)  # pytype: disable=bad-return-type  # jax-ndarray
Compute entity linking loss.

  Args:
    mention_encodings: [n_mentions, hidden_size] mention encodings to be used
      for computing the loss.
    entity_embeddings: [n_entities, hidden_size] entity embeddings table.
    mention_target_ids: [n_mentions] IDs of mentions.
    mention_target_weights: [n_mentions] per-mention weight for computing loss
      and metrics.
    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').

  Returns:
    Loss, a dictionary with metrics values, per sample infomation
    (a tuple of accuracy per mention and weight per mention).
  
language/mentionmemory/modules/mention_losses.py