docs: def _attend_over_memory(self, memory):
      """Perform multiheaded attention over `memory`.
  
      Args:
        memory: Current relational memory.
  
      Returns:
        The attended-over memory.
      """
      attention_mlp = basic.BatchApply(
          mlp.MLP([self._mem_size] * self._attention_mlp_layers))
      for _ in range(self._num_blocks):
        attended_memory = self._multihead_attention(memory)
  
        # Add a skip connection to the multiheaded attention's input.
        memory = basic.BatchApply(layer_norm.LayerNorm())(
            memory + attended_memory)
  
        # Add a skip connection to the attention_mlp's input.
        memory = basic.BatchApply(layer_norm.LayerNorm())(
            attention_mlp(memory) + memory)
  
      return memory