def sum_by_batch_position(mention_batch_positions: Array, values: Array,
                          batch_size: int) -> Array:
  """Sum values per position in the batch."""
  position2item = (
      jnp.expand_dims(jnp.arange(batch_size), 1) == mention_batch_positions)
  position2item = position2item.astype(jnp.int32)

  # [batch_size, ...]
  values_sum_per_batch_position = jnp.einsum('bm,m...->b...', position2item,
                                             values)
  return values_sum_per_batch_position
Sum values per position in the batch.
language/mentionmemory/utils/mention_utils.py