def add_memories(self, batch: Dict[str, Array], predictions: Dict[str,
                                                                    Array]):
    """Save generated memories in-memory storage."""
    mention_mask = batch['mention_target_weights'] > 0
    memory_index_end = min(self.num_total_memories,
                           self.memory_index + mention_mask.sum())
    memory_index_len = memory_index_end - self.memory_index
    indices = self.memory_permutation[self.memory_index:memory_index_end]

    # We might not save mention encodings for all target mentions.
    # First, some of them might are pad or are too close to a passage boundary
    # (in these cases we assume that `mention_target_weights` = 0).
    # Second, there might be more mentions then we actually need
    # since we limit the number of the total memories by `num_total_memories`.
    # Therefore, we create `mention_index` to select a subset of mentions,
    # which encodings we are planning to save.
    mention_index_0, mention_index_1 = jnp.nonzero(mention_mask)
    mention_index_0 = mention_index_0[:memory_index_len]
    mention_index_1 = mention_index_1[:memory_index_len]
    mention_index = (mention_index_0, mention_index_1)

    self.memory_embeddings[indices] = predictions['values'][mention_index]
    if self.memory_key_embeddings is not None:
      self.memory_key_embeddings[indices] = predictions['keys'][mention_index]
    self.memory_labels[indices] = batch['mention_target_ids'][mention_index]
    self.memory_text_hashes[indices] = batch['target_text_identifiers'][
        mention_index]
    self.memory_mention_hashes[indices] = batch['target_mention_hashes'][
        mention_index]

    # Convert to global batch positions
    n_devices, batch_size, _ = batch['text_ids'].shape
    mention_target_batch_positions = batch['mention_target_batch_positions']
    mention_target_batch_positions = (
        mention_target_batch_positions +
        np.expand_dims(np.arange(n_devices), 1) * batch_size)
    mention_target_batch_positions = mention_target_batch_positions[
        mention_index]

    self.text_entities[indices] = batch['unique_mention_ids'].reshape(
        n_devices * batch_size, -1)[mention_target_batch_positions]
    self.text_ids[indices] = batch['text_ids'].reshape(
        n_devices * batch_size, -1)[mention_target_batch_positions]
    self.start_end_positions[
        indices, 0] = batch['mention_target_start_positions'][mention_index]
    self.start_end_positions[
        indices, 1] = batch['mention_target_end_positions'][mention_index]
    self.memory_index = memory_index_end
Save generated memories in-memory storage.
language/mentionmemory/tasks/memory_generation_task.py