def __call__(self, batch: Dict[str, Any],
               auxiliary_output: Dict[str, Any]) -> Dict[str, Any]:
    """Produces memory texts and mention positions given memory IDs."""
    # We want to load memory lazily, so we do that here and not in the
    # constructor. First, we check if memory components have been loaded yet
    # (`self.memory_entity_id` is not None) and actually load them only if not.
    self.maybe_load_memory()

    def process_retrievals(prefix=''):
      top_entity_ids = auxiliary_output[prefix + 'top_entity_ids']
      top_memory_ids = auxiliary_output[prefix + 'top_memory_ids']
      if self.save_k_retrieval is not None:
        top_entity_ids = top_entity_ids[:, :, :self.save_k_retrieval]
        top_memory_ids = top_memory_ids[:, :, :self.save_k_retrieval]
      n_devices = top_entity_ids.shape[0]
      n_mentions = top_entity_ids.shape[1]
      n_retrievals = top_entity_ids.shape[2]
      assert top_entity_ids.shape == top_memory_ids.shape
      logging.info(
          'Saving %sretrievals: n_devices=%d, n_mentions=%d, n_retrievals=%d',
          prefix, n_devices, n_mentions, n_retrievals)

      features = {
          'memory_text': [],
          'memory_positions': [],
          'memory_entity_id': [],
      }
      for device_index in range(n_devices):
        memory_text_per_device = []
        memory_positions_per_device = []
        memory_entity_id_per_device = []
        for mention_index in range(n_mentions):
          memory_text_per_mention = []
          memory_positions_per_mention = []
          memory_entity_id_per_mention = []
          for r_index in range(n_retrievals):
            memory_index = top_memory_ids[device_index, mention_index, r_index]
            memory_text_per_mention.append(
                self.memory_text[memory_index].tolist())
            memory_positions_per_mention.append(
                self.memory_positions[memory_index].tolist())
            memory_entity_id_per_mention.append(
                self.memory_entity_id[memory_index].tolist())
          memory_text_per_device.append(memory_text_per_mention)
          memory_positions_per_device.append(memory_positions_per_mention)
          memory_entity_id_per_device.append(memory_entity_id_per_mention)
        features['memory_text'].append(memory_text_per_device)
        features['memory_positions'].append(memory_positions_per_device)
        features['memory_entity_id'].append(memory_entity_id_per_device)

      n_mismatch = (np.array(features['memory_entity_id']) !=
                    top_entity_ids).sum()
      if n_mismatch > 0:
        raise ValueError('Found %d mismatches amongst %d IDs in total' %
                         (n_mismatch, top_entity_ids.size))
      return {prefix + key: value for key, value in features.items()}

    retrieval_features = process_retrievals()
    if 'second_top_entity_ids' in auxiliary_output:
      retrieval_features.update(process_retrievals('second_'))
    return retrieval_features
Produces memory texts and mention positions given memory IDs.
language/mentionmemory/encoders/mention_memory_encoder.py