class_docstr: ## Compressive Transformer Layer
  
  This is the implementation of a single compressive transformer layer
docstr: * `x` is a tensor of token level feature vectors of shape `[seq_len, batch_size, d_model]`
  * `mem` is a tensor of the past token level feature vectors (memory) of shape `[mem_len, batch_size, d_model]`
  * `c_mem` is a tensor of the compressed memory `[c_mem_len, batch_size, d_model]`
  * `mask` is a matrix of shape `[seq_len, c_mem_len + mem_len + seq_len, batch_size]` or `[seq_len, c_mem_len + mem_len + seq_len, 1]`.
  `mask[i, j]` is  true if token at `i` can see token at `j`.
docstr_tok: ["*","`","x","`","is","a","tensor","of","token","level","feature","vectors","of","shape","`","[","seq_len",",","batch_size",",","d_model","]","`","*","`","mem","`","is","a","tensor","of","the","past","token","level","feature","vectors","(","memory",")","of","shape","`","[","mem_len",",","batch_size",",","d_model","]","`","*","`","c_mem","`","is","a","tensor","of","the","compressed","memory","`","[","c_mem_len",",","batch_size",",","d_model","]","`","*","`","mask","`","is","a","matrix","of","shape","`","[","seq_len",",","c_mem_len","+","mem_len","+","seq_len",",","batch_size","]","`","or","`","[","seq_len",",","c_mem_len","+","mem_len","+","seq_len",",","1","]","`",".","`","mask","[","i",",","j","]","`","is","true","if","token","at","`","i","`","can","see","token","at","`","j","`","."]
code: def forward(
      self,
      *,
      x: torch.Tensor,
      mem: Optional[torch.Tensor],
      c_mem: Optional[torch.Tensor],
      mask: torch.Tensor,
  ):
      """
      * `x` is a tensor of token level feature vectors of shape `[seq_len, batch_size, d_model]`
      * `mem` is a tensor of the past token level feature vectors (memory) of shape `[mem_len, batch_size, d_model]`
      * `c_mem` is a tensor of the compressed memory `[c_mem_len, batch_size, d_model]`
      * `mask` is a matrix of shape `[seq_len, c_mem_len + mem_len + seq_len, batch_size]` or `[seq_len, c_mem_len + mem_len + seq_len, 1]`.
      `mask[i, j]` is  true if token at `i` can see token at `j`.
      """
  
      # Normalize the vectors before doing self attention
      z = self.norm_self_attn(x)
      # Normalize and concatenate memory and compressed memory
      m_z = self.concat_memory(z, mem, c_mem)
      # Attention
      self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)
      # Add the attention results
      x = x + self.dropout(self_attn)
  
      # Normalize for feed-forward
      z = self.norm_ff(x)
      # Pass through the feed-forward network
      ff = self.feed_forward(z)
      # Add the feed-forward results back
      x = x + self.dropout(ff)
  
      #
      return x
  
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 3 for calling the CompressiveTransformerLayer obj's underlying member method code with example usage: obj.forward() and returns: x