def __call__(
      self,
      encoded_input: Array,
      mention_batch_positions: Array,
      mention_start_positions: Array,
      mention_end_positions: Array,
      mention_mask: Array,
      memory_keys: Array,
      memory_identifiers: Array,
      memory_entity_ids: Array,
      deterministic: bool,
      memory_values: Optional[Array] = None,
      text_identifiers: Optional[Array] = None,
      memory_text_entities: Optional[Array] = None,
      same_passage_memory_policy: str = 'disallow',
  ) -> Tuple[Array, Dict[str, Array], Dict[str, Array]]:
    """Perform attention update over memory table.

    Args:
      encoded_input: [batch_size, n_tokens, hidden_size] input representation.
      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.
      memory_keys: [rows, values per row, key_dim] mention memory keys. The
        number of rows in the memory table governs the recall vs speed of the
        topk similarity search. Search is performed by taking max over each row,
        and then top-k between rows. Distributing the same values over more rows
        leads to higher recall but slower search.
      memory_identifiers: [memory_size] identifier for memory vectors.
      memory_entity_ids: [memory_size] entity ids for mentions in memory table
      deterministic: don't apply dropout if true.
      memory_values: [values, memory_dim] if separate keys and values.
      text_identifiers: [n_mentions] search will not retrieve memory vectors
        with the same identifier as passage mention.
      memory_text_entities: [n_mentions, n_memory_text_entities] entity ids for
        passages where memories are coming from.
      same_passage_memory_policy: how to treat mentions from the same passage.
        Possible options: `allow`, `disallow` and `only`.

    Returns:
      Updated input, loss and logging helper dicts.
    """
    _assert_array_is_integer_or_none(mention_batch_positions,
                                     'mention_batch_positions')
    _assert_array_is_integer_or_none(mention_start_positions,
                                     'mention_start_positions')
    _assert_array_is_integer_or_none(mention_end_positions,
                                     'mention_end_positions')
    _assert_array_is_integer_or_none(memory_entity_ids, 'memory_entity_ids')
    _assert_array_is_integer_or_none(memory_identifiers, 'memory_identifiers')
    _assert_array_is_integer_or_none(memory_text_entities,
                                     'memory_text_entities')
    _assert_array_is_integer_or_none(text_identifiers, 'text_identifiers')

    loss_helpers, logging_helpers = {}, {}

    # We generate mention representations to use as queries for similarity
    # search by concatenating start and end tokens for each mention and
    # projecting the concatenation with a dense layer.
    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))

    queries = self.query_projector(
        jnp.concatenate((mention_start_encodings, mention_end_encodings),
                        axis=-1))

    loss_helpers['memory_attention_mention_encodings'] = queries

    retrieval_result = self.memory_retrieval_layer(
        queries=queries,
        memory_keys=memory_keys,
        memory_identifiers=memory_identifiers,
        memory_entity_ids=memory_entity_ids,
        memory_values=memory_values,
        text_identifiers=text_identifiers,
        memory_text_entities=memory_text_entities,
        same_passage_memory_policy=same_passage_memory_policy,
    )

    # Most of the information from retrieval_result goes to `loss_helpers`
    # except `n_disallowed`. In future, we might join these two into a single
    # dictionary.
    loss_helpers.update(retrieval_result)
    if 'n_disallowed' in retrieval_result:
      logging_helpers['n_disallowed'] = retrieval_result['n_disallowed']

    encoded_input = self.update_layer(
        encoded_input=encoded_input,
        retrieval_values=retrieval_result['top_values'],
        retrieval_scores=retrieval_result['memory_attention_weights'],
        mention_batch_positions=mention_batch_positions,
        mention_start_positions=mention_start_positions,
        mention_end_positions=mention_end_positions,
        mention_mask=mention_mask,
        deterministic=deterministic,
    )

    return encoded_input, loss_helpers, logging_helpers
Perform attention update over memory table.

    Args:
      encoded_input: [batch_size, n_tokens, hidden_size] input representation.
      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.
      memory_keys: [rows, values per row, key_dim] mention memory keys. The
        number of rows in the memory table governs the recall vs speed of the
        topk similarity search. Search is performed by taking max over each row,
        and then top-k between rows. Distributing the same values over more rows
        leads to higher recall but slower search.
      memory_identifiers: [memory_size] identifier for memory vectors.
      memory_entity_ids: [memory_size] entity ids for mentions in memory table
      deterministic: don't apply dropout if true.
      memory_values: [values, memory_dim] if separate keys and values.
      text_identifiers: [n_mentions] search will not retrieve memory vectors
        with the same identifier as passage mention.
      memory_text_entities: [n_mentions, n_memory_text_entities] entity ids for
        passages where memories are coming from.
      same_passage_memory_policy: how to treat mentions from the same passage.
        Possible options: `allow`, `disallow` and `only`.

    Returns:
      Updated input, loss and logging helper dicts.
    
language/mentionmemory/modules/memory_attention_layer.py