def __call__(
      self,
      encoded_input: Array,
      retrieval_values: Array,
      retrieval_scores: Array,
      mention_batch_positions: Array,
      mention_start_positions: Array,
      mention_end_positions: Array,
      mention_mask: Array,
      deterministic: bool,
  ) -> Array:
    """Incorporates retrieval into input.

    Args:
      encoded_input: [batch_size, n_tokens, hidden_size] input representation.
      retrieval_values: [n_mentions, n_retrieval, retrieval_dim] retrievals.
      retrieval_scores: [n_mentions, n_retrieval] retrieval scores.
      mention_batch_positions: [n_mentions] mention sample position in batch.
      mention_start_positions: [n_mentions] mention start position in input.
      mention_end_positions: [n_mentions] mention end position in input.
      mention_mask: [n_mentions] attention mask to prevent updates from padding.
      deterministic: don't apply dropout if true.

    Returns:
      Updated input.
    """  # pytype: disable=bad-return-type  # jax-ndarray
Incorporates retrieval into input.

    Args:
      encoded_input: [batch_size, n_tokens, hidden_size] input representation.
      retrieval_values: [n_mentions, n_retrieval, retrieval_dim] retrievals.
      retrieval_scores: [n_mentions, n_retrieval] retrieval scores.
      mention_batch_positions: [n_mentions] mention sample position in batch.
      mention_start_positions: [n_mentions] mention start position in input.
      mention_end_positions: [n_mentions] mention end position in input.
      mention_mask: [n_mentions] attention mask to prevent updates from padding.
      deterministic: don't apply dropout if true.

    Returns:
      Updated input.
    
language/mentionmemory/modules/retrieval_update_layers.py