class_docstr: ## Retro Model
  
  This is the Retro decoder
class_docstr_tok: ["#","#","Retro","Model","This","is","the","Retro","decoder"]
docstr: * `v_vocab` is the number of tokens in the vocabulary
  * `d_model` is the number of features in embeddings
  * `n_layers` is the number of layers in the decoder $L$
  * `ca_layers` are the layers with cross attention $P$
  * `chunk_len` is the length of a chunk
  * `n_heads` is the number of heads in attention layers
  * `d_k` is the size of attention heads
  * `d_ff` is the size of the feed-forward networks hidden layers
  * `encoder` is the nearest neighbor encoder
docstr_tok: ["*","`","v_vocab","`","is","the","number","of","tokens","in","the","vocabulary","*","`","d_model","`","is","the","number","of","features","in","embeddings","*","`","n_layers","`","is","the","number","of","layers","in","the","decoder","$","L","$","*","`","ca_layers","`","are","the","layers","with","cross","attention","$","P","$","*","`","chunk_len","`","is","the","length","of","a","chunk","*","`","n_heads","`","is","the","number","of","heads","in","attention","layers","*","`","d_k","`","is","the","size","of","attention","heads","*","`","d_ff","`","is","the","size","of","the","feed-forward","networks","hidden","layers","*","`","encoder","`","is","the","nearest","neighbor","encoder"]
code: def __init__(
      self,
      n_vocab: int,
      d_model: int,
      n_layers: int,
      ca_layers: Set[int],
      chunk_len: int,
      n_heads: int,
      d_k: int,
      d_ff: int,
      encoder: NearestNeighborEncoder,
  ):
      """
      * `v_vocab` is the number of tokens in the vocabulary
      * `d_model` is the number of features in embeddings
      * `n_layers` is the number of layers in the decoder $L$
      * `ca_layers` are the layers with cross attention $P$
      * `chunk_len` is the length of a chunk
      * `n_heads` is the number of heads in attention layers
      * `d_k` is the size of attention heads
      * `d_ff` is the size of the feed-forward networks hidden layers
      * `encoder` is the nearest neighbor encoder
      """
      super().__init__()
  
      self.ca_layers = ca_layers
      self.encoder = encoder
  
      # Token embedding layer
      self.emb = nn.Embedding(n_vocab, d_model)
      # Chunked cross attention layers $\text{C\small{CA}}$
      self.cca = nn.ModuleList(
          [
              ChunkedCrossAttention(d_model, n_heads, d_k, chunk_len)
              for _ in range(len(ca_layers))
          ]
      )
      # Attention layers $\text{A\small{TTN}}$
      self.attn = nn.ModuleList(
          [
              SelfAttention(d_model, n_heads, d_k, is_causal=True)
              for _ in range(n_layers)
          ]
      )
      # Feed forward layers $\text{F\small{FW}}$
      self.ffw = nn.ModuleList(
          [FeedForward(d_model, d_ff) for _ in range(n_layers)]
      )
      # Readout layer $\text{R\small{EAD}}$
      self.read = nn.Linear(d_model, n_vocab)
  
      # Pre-normalization layer for nearest neighbor embeddings from
      # $\text{E\small{NCODER}}(\text{R\small{ET}}(C_u)_{1 \le u \le l}, H)$
      self.norm_e = nn.LayerNorm(d_model)
  
imports: import math
  from typing import Set
  import torch
  from torch import nn
  from labml.logger import inspect
label_desc: Use image node_id 1 to create a new RetroModel object from inherited base classes: nn with example: obj = RetroModel(n_vocab, d_model, n_layers, ca_layers, chunk_len, n_heads, d_k, d_ff, encoder)