def build_mtb_positive_negative_mask(mention_batch_positions: Array,
                                     all_mention_batch_positions: Array,
                                     mention_target_ids: Array,
                                     all_mention_target_ids: Array,
                                     mtb_positive_mask: Array):
  """Decides which samples are positive / negative for MTB loss.

  On of the argument `mtb_positive_mask` contains information regarding how
  many unique common entities does different passages in the batch contain.
  mtb_positive_mask[i, j] is True <=> a passage in the local batch that contains
  mention i has enough (typically, at least 2) unique common entities with
  a passage in the global batch that corresponds to mention j.

  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 MTB 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
  (5) mtb_positive_mask[i, j] is True

  hard_negative_mask[i, j] is True <=> mention j in `all_mention_target_ids`
  could be used as a hard negative sample for mention i in `mention_target_ids`
  for the MTB 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
  (5) mtb_positive_mask[i, j] is False

  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 MTB 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
  OR hard_negative_mask[i, j] is True

  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.
    mtb_positive_mask: shape [n_mentions, n_devices * n_mentions]. Whether
      mentions are considered as positive in MTB sense -- their passage share at
      least 2 unique common entities.

  Returns:
    Positive_mask, hard_negative_mask and negative_mask.
  """
  same_passage = mention_utils.all_compare(mention_batch_positions,
                                           all_mention_batch_positions)
  not_same_passage_mask = jnp.logical_not(same_passage)
  same_entity_mask = mention_utils.all_compare(mention_target_ids,
                                               all_mention_target_ids)
  not_same_entity_mask = jnp.logical_not(same_entity_mask)
  not_pad = jnp.expand_dims(mention_target_ids, 1) > 0
  not_pad_all = 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)
  positive_mask = jnp.logical_and(positive_mask, not_pad_all)

  hard_negative_mask = jnp.logical_and(positive_mask,
                                       jnp.logical_not(mtb_positive_mask))
  positive_mask = jnp.logical_and(positive_mask, mtb_positive_mask)

  negative_mask = jnp.logical_and(not_same_entity_mask, not_same_passage_mask)
  negative_mask = jnp.logical_and(negative_mask, not_pad)
  negative_mask = jnp.logical_and(negative_mask, not_pad_all)
  negative_mask = jnp.logical_or(negative_mask, hard_negative_mask)
  return positive_mask, hard_negative_mask, negative_mask
Decides which samples are positive / negative for MTB loss.

  On of the argument `mtb_positive_mask` contains information regarding how
  many unique common entities does different passages in the batch contain.
  mtb_positive_mask[i, j] is True <=> a passage in the local batch that contains
  mention i has enough (typically, at least 2) unique common entities with
  a passage in the global batch that corresponds to mention j.

  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 MTB 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
  (5) mtb_positive_mask[i, j] is True

  hard_negative_mask[i, j] is True <=> mention j in `all_mention_target_ids`
  could be used as a hard negative sample for mention i in `mention_target_ids`
  for the MTB 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
  (5) mtb_positive_mask[i, j] is False

  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 MTB 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
  OR hard_negative_mask[i, j] is True

  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.
    mtb_positive_mask: shape [n_mentions, n_devices * n_mentions]. Whether
      mentions are considered as positive in MTB sense -- their passage share at
      least 2 unique common entities.

  Returns:
    Positive_mask, hard_negative_mask and negative_mask.
  
language/mentionmemory/modules/mention_losses.py