def __call__(
      self,
      encoded_input: Array,
      mention_batch_positions: Array,
      mention_start_positions: Array,
      mention_end_positions: Array,
      mention_mask: Array,
      entity_embeddings: Array,
  ) -> Dict[str, Array]:
    """Perform attention update over entity embedding table.

    Args:
      encoded_input: [batch_size, n_tokens, hidden_size].
      mention_batch_positions: [n_mentions].
      mention_start_positions: [n_mentions].
      mention_end_positions: [n_mentions].
      mention_mask: [n_mentions] attention mask to prevent updates from padding
        mentions.
      entity_embeddings: entity embedding table.

    Returns:
      Updated input, mention encodings and entity attention scores.
    """

    mention_start_encodings = jut.matmul_2d_index_select(
        encoded_input, (mention_batch_positions, mention_start_positions))
    mention_end_encodings = jut.matmul_2d_index_select(
        encoded_input, (mention_batch_positions, mention_end_positions))
    mention_encodings = self.mention_query_projector(
        jnp.concatenate((mention_start_encodings, mention_end_encodings),
                        axis=-1))

    scores = jnp.einsum('qd,ed->qe', mention_encodings, entity_embeddings)
    attention_weights = nn.softmax(scores, axis=-1)

    retrieved_values = jnp.einsum('qe,ed->qd', attention_weights,
                                  entity_embeddings)
    retrieved_values = self.entity_projector(retrieved_values)
    retrieved_values = retrieved_values * jnp.expand_dims(mention_mask, -1)

    encoded_input = jut.matmul_2d_index_add(
        encoded_input, (mention_batch_positions, mention_start_positions),
        retrieved_values)
    encoded_input = self.layer_norm(encoded_input)

    # The cosine similarity is computed as dot product divided by norms of
    # both vectors.
    mention_encodings_norm = jnp.linalg.norm(mention_encodings, axis=-1)
    entity_embeddings_norm = jnp.linalg.norm(entity_embeddings, axis=-1)
    cos_scores = scores
    cos_scores = cos_scores / (
        _SMALL_NUMBER + jnp.expand_dims(mention_encodings_norm, 1))
    cos_scores = cos_scores / (
        _SMALL_NUMBER + jnp.expand_dims(entity_embeddings_norm, 0))

    return {
        'encoded_output': encoded_input,
        'mention_encodings': mention_encodings,
        'cosine_similarity': cos_scores,
        'attention_weights': attention_weights,
    }
Perform attention update over entity embedding table.

    Args:
      encoded_input: [batch_size, n_tokens, hidden_size].
      mention_batch_positions: [n_mentions].
      mention_start_positions: [n_mentions].
      mention_end_positions: [n_mentions].
      mention_mask: [n_mentions] attention mask to prevent updates from padding
        mentions.
      entity_embeddings: entity embedding table.

    Returns:
      Updated input, mention encodings and entity attention scores.
    
language/mentionmemory/modules/entity_attention_layer.py