def get_predictions_sum(attention_weights: Array, memory_entity_ids: Array,
                        weights: Array, entity_vocab_size: int) -> Array:
  """Get entity with the largest sum of attention weights over its mentions."""
  attention_weights = jax.lax.stop_gradient(attention_weights)
  memory_entity_ids = jax.lax.stop_gradient(memory_entity_ids)
  weights = jax.lax.stop_gradient(weights)

  n_mentions = attention_weights.shape[0]
  attention_weights_per_entity = jnp.zeros((n_mentions, entity_vocab_size),
                                           dtype=attention_weights.dtype)
  attention_weights_per_entity = attention_weights_per_entity.at[
      jnp.expand_dims(jnp.arange(n_mentions), 1),
      memory_entity_ids].add(attention_weights)
  predictions = jnp.argmax(attention_weights_per_entity, axis=1)
  predictions = predictions * weights
  return predictions
Get entity with the largest sum of attention weights over its mentions.
language/mentionmemory/tasks/mention_based_entity_qa_task.py