def get_num_common_unique_items(batch_positions: Array, ids: Array,
                                batch_size: int) -> Array:
  """Get the number of unique entity IDs shared between samples in the batch.

  The function produces two matrices `num_common_ids_between_samples` of shape
  [batch_size, batch_size * n_devices] and `num_common_ids_between_mentions` of
  shape [n_mentions, n_mentions * n_devices]. The first matrix is the number
  of unique entity IDs that exists in both sample A and sample B. The second
  matrix is similar, but is indexed with respect to individual mentions --
  how many unique entity IDs exists for a sample that contains mention A and
  a sample that contains mention B.

  More formally, let `all_ids` be IDs concatenated from all devices and
  `global_batch_position` be global batch positions produced by the
  `get_globally_consistent_batch_positions` method.

  `num_common_ids_between_samples[i, j] = k` <=> The number of unique IDs in the
  intersection of ids[batch_positions == i] and
  all_ids[global_batch_position == j] is k.

  `num_common_ids_between_mentions[i, j] = k` <=> If b_i is a batch position
  corresponding to a local mention i and b_j is a global batch position
  corresponding to a global mention j then
  num_common_ids_between_samples[b_i, b_j] = k

  Args:
    batch_positions: position of a mention in its batch.
    ids: IDs of mentions in the batch.
    batch_size: batch size per device.

  Returns:
    Matrices with the number of unique common elements between a row in local
    batch and a row in global batch. Shapes are
    [batch_size, n_devices * batch_size] between samples in a batch and
    [n_mentions, n_devices * n_mentions] between mentions in a batch.
  """
  ids = mask_duplicate_ids(batch_positions, ids)

  device_id = get_device_id('batch')
  if device_id is not None:
    all_ids = jax.lax.all_gather(ids, 'batch')
    n_devices = all_ids.shape[0]
    all_ids = all_ids.reshape([-1])
  else:
    n_devices = 1
    all_ids = ids

  # [n_mentions, n_global_mentions]
  same_ids = all_compare_without_pad(ids, all_ids)
  same_ids = same_ids.astype(jnp.int32)

  (_, global_batch_positions) = get_globally_consistent_batch_positions(
      batch_positions, batch_size)

  # [batch_size, n_mentions]
  position2item = (
      jnp.expand_dims(jnp.arange(batch_size), 1) == batch_positions)
  position2item = position2item.astype(jnp.int32)

  # [global_batch_size, n_global_mentions]
  item2global_position = (
      jnp.expand_dims(jnp.arange(n_devices * batch_size),
                      1) == global_batch_positions)
  item2global_position = item2global_position.astype(jnp.int32)

  # [batch_size, n_global_mentions]
  num_common_ids_between_samples = sum_by_batch_position(
      batch_positions, same_ids, batch_size)

  # [batch_size, global_batch_size]
  num_common_ids_between_samples = jnp.einsum('bn,gn->bg',
                                              num_common_ids_between_samples,
                                              item2global_position)
  num_common_ids_between_mentions = jnp.einsum('bm,bg->mg', position2item,
                                               num_common_ids_between_samples)
  num_common_ids_between_mentions = jnp.einsum('gn,mg->mn',
                                               item2global_position,
                                               num_common_ids_between_mentions)
  return num_common_ids_between_samples, num_common_ids_between_mentions  # pytype: disable=bad-return-type  # jax-ndarray
Get the number of unique entity IDs shared between samples in the batch.

  The function produces two matrices `num_common_ids_between_samples` of shape
  [batch_size, batch_size * n_devices] and `num_common_ids_between_mentions` of
  shape [n_mentions, n_mentions * n_devices]. The first matrix is the number
  of unique entity IDs that exists in both sample A and sample B. The second
  matrix is similar, but is indexed with respect to individual mentions --
  how many unique entity IDs exists for a sample that contains mention A and
  a sample that contains mention B.

  More formally, let `all_ids` be IDs concatenated from all devices and
  `global_batch_position` be global batch positions produced by the
  `get_globally_consistent_batch_positions` method.

  `num_common_ids_between_samples[i, j] = k` <=> The number of unique IDs in the
  intersection of ids[batch_positions == i] and
  all_ids[global_batch_position == j] is k.

  `num_common_ids_between_mentions[i, j] = k` <=> If b_i is a batch position
  corresponding to a local mention i and b_j is a global batch position
  corresponding to a global mention j then
  num_common_ids_between_samples[b_i, b_j] = k

  Args:
    batch_positions: position of a mention in its batch.
    ids: IDs of mentions in the batch.
    batch_size: batch size per device.

  Returns:
    Matrices with the number of unique common elements between a row in local
    batch and a row in global batch. Shapes are
    [batch_size, n_devices * batch_size] between samples in a batch and
    [n_mentions, n_devices * n_mentions] between mentions in a batch.
  
language/mentionmemory/utils/mention_utils.py