def get_all_encodings_and_target_ids(
    mention_encodings: Array, mention_target_ids: Array) -> Tuple[Array, Array]:
  """Gather encodings and IDs from all devices."""
  hidden_dim = mention_encodings.shape[-1]
  device_id = mention_utils.get_device_id('batch')
  if device_id is not None:
    all_mention_encodings = jax.lax.all_gather(mention_encodings, 'batch')
    all_mention_encodings = all_mention_encodings.reshape([-1, hidden_dim])

    all_mention_target_ids = jax.lax.all_gather(mention_target_ids, 'batch')
    all_mention_target_ids = all_mention_target_ids.reshape([-1])
  else:
    all_mention_encodings = mention_encodings
    all_mention_target_ids = mention_target_ids

  return all_mention_encodings, all_mention_target_ids
Gather encodings and IDs from all devices.
language/mentionmemory/modules/mention_losses.py