class_docstr: ## Attention Reconstruction Loss
  
  Attention reconstruction loss recreates the self-attention output with
  uncompressed memory and with compressed memory and calculates the mean squared error
  between the two. It does this without positional encoding.
  
  When calculating and training the compression function $f_c$ with attention
  reconstruction loss, all parameters but $f_c$ are frozen.
  This includes key/value projections and bias/scaling after normalization.
  
  Since this loss can be computed independently of the cross-entropy-loss of the model
  you can have a separate optimizer that only updates $f_c$.
  However, we use the same optimizer to update $f_c$ so when calculating
  attention reconstruction loss, we detach all other parameters except $f_c$
  from the gradient computation.
code: def __call__(self, h: List[torch.Tensor], mem: List[torch.Tensor]):
      # Calculate the losses for each layer
      losses = [
          self.calc_loss(layer, h[n], mem[n])
          for n, layer in enumerate(self.layers)
      ]
      # Sum of the losses
      return sum(losses)
  
imports: from typing import Optional, List
  import torch
  import torch.nn.functional
  from torch import nn
  from labml_helpers.module import Module, TypedModuleList
  from labml_nn.transformers.feed_forward import FeedForward
  from labml_nn.transformers.mha import PrepareForMultiHeadAttention
  from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
  from labml_nn.utils import clone_module_list
label_desc: Use image node_id 6 for calling the AttentionReconstructionLoss obj's underlying member method code with example usage: obj.__call__(h, mem) and returns: sum