def get_predictions_max(attention_weights: Array, memory_entity_ids: Array,
                        weights: Array) -> Array:
  """Predict entity ID based on a memory with the largest attention weight."""
  # `stop_gradient` is a safety check so the model doesn't keep activations
  # around, which could be expensive.
  attention_weights = jax.lax.stop_gradient(attention_weights)
  memory_entity_ids = jax.lax.stop_gradient(memory_entity_ids)
  weights = jax.lax.stop_gradient(weights)
  memory_with_largest_attn = jnp.argmax(attention_weights, axis=1)
  predictions = jnp.take_along_axis(
      memory_entity_ids, jnp.expand_dims(memory_with_largest_attn, 1), axis=1)
  predictions = jnp.squeeze(predictions, axis=1)
  predictions = predictions * weights
  return predictions
Predict entity ID based on a memory with the largest attention weight.
language/mentionmemory/tasks/mention_based_entity_qa_task.py