def get_batch_and_retrievals_entity_overlap(
    mention_target_batch_positions: Array,
    mention_target_ids: Array,
    mention_target_weights: Array,
    memory_text_entities: Array,
    batch_size: int,
):
  """Compute the overlap between entities in the batch and in retrievals.

  Args:
    mention_target_batch_positions: [n_target_mentions] position of a mention in
      its batch.
    mention_target_ids: [n_target_mentions] IDs of mentions.
    mention_target_weights: [n_target_mentions] per-mention weight for computing
      loss and metrics.
    memory_text_entities: [n_retrievals, n_memory_text_entities] IDs of mentions
      in the passage where memory is coming from. Note, entities in the same
      retrieval are assumed to be unique.
    batch_size: batch size.

  Returns:
    Array of shape [batch_size, n_retrievals] with the number of
      overlapping unique entity IDs in the batch and in the retrieval results.
  """
  n_target_mentions = mention_target_batch_positions.shape[0]
  n_retrievals = memory_text_entities.shape[0]
  n_memory_text_entities = memory_text_entities.shape[1]

  # Step 1: de-duplicate entities in the batch.
  # [n_target_mentions]
  mention_target_ids = mention_target_ids * mention_target_weights
  # [n_target_mentions]
  batch_ids = mention_utils.mask_duplicate_ids(mention_target_batch_positions,
                                               mention_target_ids)

  # Step 2: compare all entities in the batch against all entities in the
  # retrieval result.
  memory_text_entities = memory_text_entities.reshape(
      [n_retrievals * n_memory_text_entities])
  # [n_target_mentions, n_mentions * k_top * n_memory_text_entities]
  mention_id_in_retrieved_passages = mention_utils.all_compare_without_pad(
      batch_ids, memory_text_entities)
  mention_id_in_retrieved_passages = mention_id_in_retrieved_passages.astype(
      jnp.int32)

  # Step 3: sum up the comparison results by retrieval ID
  # [n_target_mentions, n_retrievals]
  mention_id_in_retrieved_passages = mention_id_in_retrieved_passages.reshape(
      [n_target_mentions, n_retrievals, n_memory_text_entities]).sum(-1)

  # Step 4: sum up the comparison results by batch position
  # [batch_size, n_retrievals]
  num_common_ids_between_samples = mention_utils.sum_by_batch_position(
      mention_target_batch_positions, mention_id_in_retrieved_passages,
      batch_size)

  return num_common_ids_between_samples
Compute the overlap between entities in the batch and in retrievals.

  Args:
    mention_target_batch_positions: [n_target_mentions] position of a mention in
      its batch.
    mention_target_ids: [n_target_mentions] IDs of mentions.
    mention_target_weights: [n_target_mentions] per-mention weight for computing
      loss and metrics.
    memory_text_entities: [n_retrievals, n_memory_text_entities] IDs of mentions
      in the passage where memory is coming from. Note, entities in the same
      retrieval are assumed to be unique.
    batch_size: batch size.

  Returns:
    Array of shape [batch_size, n_retrievals] with the number of
      overlapping unique entity IDs in the batch and in the retrieval results.
  
language/mentionmemory/modules/mention_losses.py