def build_coref_positive_negative_mask(mention_batch_positions: Array,
                                       all_mention_batch_positions: Array,
                                       mention_target_ids: Array,
                                       all_mention_target_ids: Array):
  """Decides which samples are positive / negative for coref resolution loss.

  positive_mask[i, j] is True <=> mention j in `all_mention_target_ids`
  could be used as a positive sample for mention i in `mention_target_ids`
  for the coreference resolution loss. It requires the following conditions:
  (1) mention_target_ids[i] == all_mention_target_ids[j]
  (2) mention i and mention j are not from the same passage
  (3) mention_target_ids[i] != 0 <=> is not padding
  (4) all_mention_target_ids[j] != 0 <=> is not padding

  negative_mask[i, j] is True <=> mention j in `all_mention_target_ids`
  could be used as a negative sample for mention i in `mention_target_ids`
  for the coreference resolution loss. It requires the following conditions:
  (1) mention_target_ids[i] != all_mention_target_ids[j]
  (2) mention i and mention j are not from the same passage
  (3) mention_target_ids[i] != 0 <=> is not padding
  (4) all_mention_target_ids[j] != 0 <=> is not padding

  Args:
    mention_batch_positions: position of a mention in its batch.
    all_mention_batch_positions: position of a global mention in its batch.
    mention_target_ids: shape [n_mentions]. IDs of mentions on the current
      device.
    all_mention_target_ids: shape [n_mentions]. IDs of all mentions on all
      devices.

  Returns:
    Positive and negative masks.
  """
  same_passage_mask = mention_utils.all_compare(mention_batch_positions,
                                                all_mention_batch_positions)
  not_same_passage_mask = jnp.logical_not(same_passage_mask)
  same_entity_mask = mention_utils.all_compare(mention_target_ids,
                                               all_mention_target_ids)
  not_same_entity_mask = jnp.logical_not(same_entity_mask)
  # Mask rows where base mention is PAD. In other words, none of the mentions
  # are positive or negative examples for PAD mentions.
  not_pad_mask = jnp.expand_dims(mention_target_ids, 1) > 0
  # Mask columns where target mention is PAD. In other words, we don't consider
  # PAD mentions to be either positive nor negative example.
  not_pad_all_mask = jnp.expand_dims(all_mention_target_ids, 0) > 0

  positive_mask = jnp.logical_and(same_entity_mask, not_same_passage_mask)
  positive_mask = jnp.logical_and(positive_mask, not_pad_mask)
  positive_mask = jnp.logical_and(positive_mask, not_pad_all_mask)

  negative_mask = jnp.logical_and(not_same_entity_mask, not_same_passage_mask)
  negative_mask = jnp.logical_and(negative_mask, not_pad_mask)
  negative_mask = jnp.logical_and(negative_mask, not_pad_all_mask)
  return positive_mask, negative_mask
Decides which samples are positive / negative for coref resolution loss.

  positive_mask[i, j] is True <=> mention j in `all_mention_target_ids`
  could be used as a positive sample for mention i in `mention_target_ids`
  for the coreference resolution loss. It requires the following conditions:
  (1) mention_target_ids[i] == all_mention_target_ids[j]
  (2) mention i and mention j are not from the same passage
  (3) mention_target_ids[i] != 0 <=> is not padding
  (4) all_mention_target_ids[j] != 0 <=> is not padding

  negative_mask[i, j] is True <=> mention j in `all_mention_target_ids`
  could be used as a negative sample for mention i in `mention_target_ids`
  for the coreference resolution loss. It requires the following conditions:
  (1) mention_target_ids[i] != all_mention_target_ids[j]
  (2) mention i and mention j are not from the same passage
  (3) mention_target_ids[i] != 0 <=> is not padding
  (4) all_mention_target_ids[j] != 0 <=> is not padding

  Args:
    mention_batch_positions: position of a mention in its batch.
    all_mention_batch_positions: position of a global mention in its batch.
    mention_target_ids: shape [n_mentions]. IDs of mentions on the current
      device.
    all_mention_target_ids: shape [n_mentions]. IDs of all mentions on all
      devices.

  Returns:
    Positive and negative masks.
  
language/mentionmemory/modules/mention_losses.py