def same_entity_set_retrieval_loss(
    mention_target_batch_positions: Array,
    mention_target_ids: Array,
    mention_target_weights: Array,
    mention_batch_positions: Array,
    mention_mask: Array,
    memory_text_entities: Array,
    memory_attention_weights: Array,
    memory_mask: Array,
    batch_size: int,
    same_entity_set_target_threshold: int,
):
  """Computes same-entity-set-retrieval loss.

  We want to maximize attention scores received by memories which passages have
  large enough entity overlap (at least `same_entity_set_target_threshold`
  unique entities in common) with query mention's passage.

  User specifies which entity IDs exist in the current batch via the arguments:
  `mention_target_batch_positions`, `mention_target_ids` and
  `mention_target_weights`. Note that these correspond to linked mentions. While
  we retrieve memories for all mentions, we don't know entity IDs for non-linked
  mentions and they are not needed for computing entity overlap.

  Formally, let E(j) be the set of entity IDs in the j-th sample in the batch.
  k-th retrieval for i-th mention is "correct" if and only if
  IntersectionSize(E(mention_batch_positions[i]), memory_text_entities[i, j]) is
  at least `same_entity_set_target_threshold`.

  The loss is first computed for every mention separately and is negative log of
  sum of memory_attention_weights[i, j] where j is a "correct" retrieval for the
  i-th mention. The final loss is average of losses for those mentions which
  have at least one "correct" retrieval.

  Args:
    mention_target_batch_positions: [n_target_mentions] position of a linked
      (target) mention in its batch.
    mention_target_ids: [n_target_mentions] IDs of mentions.
    mention_target_weights: [n_target_mentions] per-mention weight for linked
      (target) mentions indicating whether a mention is padding or not.
    mention_batch_positions: [n_mentions] position of a mention in its batch.
    mention_mask: [n_mentions] whether a mention is padding or not.
    memory_text_entities: [n_mentions, k_top, n_memory_text_entities] IDs of
      mentions in the passage where memory is coming from.
    memory_attention_weights: [n_mentions, k_top] attention weights for the
      retrieval results.
    memory_mask: [n_mentions, k_top] which retrievals to use and which are to
      ignore in the loss computations. Typical usage is to ignore "disallowed"
      or same-passage retrievals.
    batch_size: batch size.
    same_entity_set_target_threshold: how many common entities needs to be
      between memory's passage and mention's passage in order to treat retrieval
      result as positive. If it's equal to 2 this loss becomes
      `same-mtb-retrieval` loss.

  Returns:
    Tuple of scalar loss, avg correct probability and normalizing factor.
  """
  n_mentions = memory_text_entities.shape[0]
  k_top = memory_text_entities.shape[1]

  # [batch_size, n_mentions * k_top]
  num_common_ids_between_samples = get_batch_and_retrievals_entity_overlap(
      mention_target_batch_positions=mention_target_batch_positions,
      mention_target_ids=mention_target_ids,
      mention_target_weights=mention_target_weights,
      memory_text_entities=memory_text_entities.reshape(
          [n_mentions * k_top, -1]),
      batch_size=batch_size,
  )

  # [batch_size, n_mentions, k_top]
  num_common_ids_between_samples = num_common_ids_between_samples.reshape(
      [batch_size, n_mentions, k_top])

  # [batch_size, n_mentions]
  position2item = (
      jnp.expand_dims(jnp.arange(batch_size), 1) == mention_batch_positions)
  position2item = position2item.astype(jnp.int32)

  # [n_mentions,  k_top]
  num_common_ids_between_mentions = jnp.einsum('bm,bmk->mk', position2item,
                                               num_common_ids_between_samples)

  # compute which retrievals have enough common elements with the
  # passage in the batch (at least `same_entity_set_target_threshold`) and
  # therefore, should be marked as "correct" retrievals.
  # [n_mentions, k_top]
  enough_common_elements = (
      num_common_ids_between_mentions >= same_entity_set_target_threshold)
  enough_common_elements = enough_common_elements.astype(jnp.int32)
  correct_retrievals_mask = enough_common_elements * memory_mask
  incorrect_retrievals_mask = (1 - enough_common_elements) * memory_mask

  # [n_mentions]
  correct_retrieval_exists = correct_retrievals_mask.sum(-1) > 0
  incorrect_retrieval_exists = incorrect_retrievals_mask.sum(-1) > 0
  loss_mask = jnp.logical_and(correct_retrieval_exists,
                              incorrect_retrieval_exists)
  loss_mask = loss_mask.astype(mention_mask.dtype) * mention_mask
  loss_mask = loss_mask.astype(jnp.float32)

  # [n_mentions, k_top]
  correct_retrievals_mask = correct_retrievals_mask.astype(jnp.float32)

  # compute loss and metrics
  # [n_mentions, k_top]
  correct_probs = jnp.einsum('mk,mk->m', correct_retrievals_mask,
                             memory_attention_weights)
  avg_probs = correct_probs * loss_mask
  loss = -jnp.log(correct_probs + _SMALL_NUMBER)
  loss = loss * loss_mask

  return loss.sum(), avg_probs.sum(), loss_mask.sum(),
Computes same-entity-set-retrieval loss.

  We want to maximize attention scores received by memories which passages have
  large enough entity overlap (at least `same_entity_set_target_threshold`
  unique entities in common) with query mention's passage.

  User specifies which entity IDs exist in the current batch via the arguments:
  `mention_target_batch_positions`, `mention_target_ids` and
  `mention_target_weights`. Note that these correspond to linked mentions. While
  we retrieve memories for all mentions, we don't know entity IDs for non-linked
  mentions and they are not needed for computing entity overlap.

  Formally, let E(j) be the set of entity IDs in the j-th sample in the batch.
  k-th retrieval for i-th mention is "correct" if and only if
  IntersectionSize(E(mention_batch_positions[i]), memory_text_entities[i, j]) is
  at least `same_entity_set_target_threshold`.

  The loss is first computed for every mention separately and is negative log of
  sum of memory_attention_weights[i, j] where j is a "correct" retrieval for the
  i-th mention. The final loss is average of losses for those mentions which
  have at least one "correct" retrieval.

  Args:
    mention_target_batch_positions: [n_target_mentions] position of a linked
      (target) mention in its batch.
    mention_target_ids: [n_target_mentions] IDs of mentions.
    mention_target_weights: [n_target_mentions] per-mention weight for linked
      (target) mentions indicating whether a mention is padding or not.
    mention_batch_positions: [n_mentions] position of a mention in its batch.
    mention_mask: [n_mentions] whether a mention is padding or not.
    memory_text_entities: [n_mentions, k_top, n_memory_text_entities] IDs of
      mentions in the passage where memory is coming from.
    memory_attention_weights: [n_mentions, k_top] attention weights for the
      retrieval results.
    memory_mask: [n_mentions, k_top] which retrievals to use and which are to
      ignore in the loss computations. Typical usage is to ignore "disallowed"
      or same-passage retrievals.
    batch_size: batch size.
    same_entity_set_target_threshold: how many common entities needs to be
      between memory's passage and mention's passage in order to treat retrieval
      result as positive. If it's equal to 2 this loss becomes
      `same-mtb-retrieval` loss.

  Returns:
    Tuple of scalar loss, avg correct probability and normalizing factor.
  
language/mentionmemory/modules/mention_losses.py