def get_globally_consistent_batch_positions(
    mention_batch_positions: Array, batch_size: int) -> Tuple[Array, Array]:
  """Adjust batch positions to be unique in the global batch.

  If the function is executed only on a single device (outside of `pmap`) then
  the results is just `mention_batch_positions`. Otherwise, the function
  returns mention_batch_positions for the first device, mention_batch_positions
  + batch size for the second device, 2*batch_size + mention_batch_positions for
  third, etc. The method returns local and global mention_batch_positions with
  adjusted batch positions.

  For example, let batch size be 3. There are two devices with
  mention_batch_positions equal to [0, 1, 2] on the first device
  and [2, 0, 0] on the second device. Then the function
  returns [0, 1, 2], [0, 1, 2, 5, 3, 3] on the first device
  and [5, 3, 3], [0, 1, 2, 5, 3, 3] on the second device.

  Args:
    mention_batch_positions: position of a mention in its batch.
    batch_size: batch size per device.

  Returns:
    Local mention_batch_positions and global mention_batch_positions with
    globally consistent IDs.
  """
  device_id = get_device_id('batch')
  if device_id is not None:
    mention_batch_positions = mention_batch_positions + batch_size * device_id
    all_mention_batch_positions = jax.lax.all_gather(mention_batch_positions,
                                                     'batch')
    all_mention_batch_positions = all_mention_batch_positions.reshape([-1])
    return mention_batch_positions, all_mention_batch_positions
  else:
    return mention_batch_positions, mention_batch_positions
Adjust batch positions to be unique in the global batch.

  If the function is executed only on a single device (outside of `pmap`) then
  the results is just `mention_batch_positions`. Otherwise, the function
  returns mention_batch_positions for the first device, mention_batch_positions
  + batch size for the second device, 2*batch_size + mention_batch_positions for
  third, etc. The method returns local and global mention_batch_positions with
  adjusted batch positions.

  For example, let batch size be 3. There are two devices with
  mention_batch_positions equal to [0, 1, 2] on the first device
  and [2, 0, 0] on the second device. Then the function
  returns [0, 1, 2], [0, 1, 2, 5, 3, 3] on the first device
  and [5, 3, 3], [0, 1, 2, 5, 3, 3] on the second device.

  Args:
    mention_batch_positions: position of a mention in its batch.
    batch_size: batch size per device.

  Returns:
    Local mention_batch_positions and global mention_batch_positions with
    globally consistent IDs.
  
language/mentionmemory/utils/mention_utils.py