def mask_duplicate_ids(batch_positions: Array, ids: Array) -> Array:
  """Zero out duplicate items within the same batch position.

  The function is a variation of a `unique` function applied to ids array
  batch_position-wise. The difference with a `np.unique` is that instead of
  discarding repeated elements (ids), the function makes them 0. For example,
  if `batch_positions` is [1, 1, 1, 2, 2] and `ids` is [1, 1, 2, 2, 2] then
  the output is [1, 0, 2, 2, 0].

  Args:
    batch_positions: position of a mention in its batch.
    ids: IDs of mentions in the batch.

  Returns:
    A modified version of `ids` where all duplicate IDs in the same batch
    position are set to zero.
  """
  same_position = all_compare(batch_positions, batch_positions)
  same_ids = all_compare(ids, ids)
  same_ids = jnp.logical_and(same_ids, same_position)
  item_is_not_duplicate = jnp.tril(same_ids).sum(axis=-1) <= 1
  ids = ids * item_is_not_duplicate
  return ids
Zero out duplicate items within the same batch position.

  The function is a variation of a `unique` function applied to ids array
  batch_position-wise. The difference with a `np.unique` is that instead of
  discarding repeated elements (ids), the function makes them 0. For example,
  if `batch_positions` is [1, 1, 1, 2, 2] and `ids` is [1, 1, 2, 2, 2] then
  the output is [1, 0, 2, 2, 0].

  Args:
    batch_positions: position of a mention in its batch.
    ids: IDs of mentions in the batch.

  Returns:
    A modified version of `ids` where all duplicate IDs in the same batch
    position are set to zero.
  
language/mentionmemory/utils/mention_utils.py