docs: def _multihead_attention(self, memory):
      """Perform multi-head attention from 'Attention is All You Need'.
  
      Implementation of the attention mechanism from
      https://arxiv.org/abs/1706.03762.
  
      Args:
        memory: Memory tensor to perform attention on.
  
      Returns:
        new_memory: New memory tensor.
      """
      key_size = self._key_size
      value_size = self._head_size
  
      qkv_size = 2 * key_size + value_size
      total_size = qkv_size * self._num_heads  # Denote as F.
      qkv = basic.BatchApply(basic.Linear(total_size))(memory)
      qkv = basic.BatchApply(layer_norm.LayerNorm())(qkv)
  
      mem_slots = memory.get_shape().as_list()[1]  # Denoted as N.
  
      # [B, N, F] -> [B, N, H, F/H]
      qkv_reshape = basic.BatchReshape([mem_slots, self._num_heads,
                                        qkv_size])(qkv)
  
      # [B, N, H, F/H] -> [B, H, N, F/H]
      qkv_transpose = tf.transpose(qkv_reshape, [0, 2, 1, 3])
      q, k, v = tf.split(qkv_transpose, [key_size, key_size, value_size], -1)
  
      q *= key_size ** -0.5
      dot_product = tf.matmul(q, k, transpose_b=True)  # [B, H, N, N]
      weights = tf.nn.softmax(dot_product)
  
      output = tf.matmul(weights, v)  # [B, H, N, V]
  
      # [B, H, N, V] -> [B, N, H, V]
      output_transpose = tf.transpose(output, [0, 2, 1, 3])
  
      # [B, N, H, V] -> [B, N, H * V]
      new_memory = basic.BatchFlatten(preserve_dims=2)(output_transpose)
      return new_memory