class_docstr: ## Retro Model
  
  This is the Retro decoder
class_docstr_tok: ["#","#","Retro","Model","This","is","the","Retro","decoder"]
docstr: * `x` is the input sequence, $X$ of shape `[batch_size, seq_len]`
  * `ret` are the retrieved neighbors
   $      ext{R\small{ET}}(C_u)_{1 \le u \le l}$
   of shape `[batch_size, chunks, neighbors, neighbor_len]`
docstr_tok: ["*","`","x","`","is","the","input","sequence",",","$","X","$","of","shape","`","[","batch_size",",","seq_len","]","`","*","`","ret","`","are","the","retrieved","neighbors","$","ext","{","R\\small","{","ET","}","}","(","C_u",")","_","{","1","\\le","u","\\le","l","}","$","of","shape","`","[","batch_size",",","chunks",",","neighbors",",","neighbor_len","]","`"]
code: def forward(self, x: torch.Tensor, ret: torch.Tensor):
      """
      * `x` is the input sequence, $X$ of shape `[batch_size, seq_len]`
      * `ret` are the retrieved neighbors
       $\text{R\small{ET}}(C_u)_{1 \le u \le l}$
       of shape `[batch_size, chunks, neighbors, neighbor_len]`
      """
  
      # Get input embeddings $H \leftarrow \text{E\small{MB}}(X)$
      h = self.emb(x)
  
      # Embeddings of the retrieved neighbors
      # $E^j_u = \text{E\small{MB}}_{\text{enc}}\big(\text{R\small{ET}}(C_u)^j\big)$.
      #
      # We use same embeddings for both input and neighbors
      ret_emb = self.emb(ret)
  
      # Keep index of the chunked cross attention layer
      p_ca = 0
      # For all layers $p \in [1, L]$
      for p in range(len(self.attn)):
          # Causal self attention $H \leftarrow \text{A\small{TTN}}(H)$
          h = self.attn[p](h)
  
          # Get encoder embeddings before the first $\text{C\small{CA}}$ layer,
          # when $p = \min(P)$
          if self.ca_layers and p == min(self.ca_layers):
              # $E = \text{E\small{NCODER}}(\text{R\small{ET}}(C_u)_{1 \le u \le l}, H)$
              #
              # We passed the embeddings of $\text{R\small{ET}}(C_u)_{1 \le u \le l}$ to encoder.
              e = self.encoder(ret_emb, h)
              # Normalize encoder embeddings
              e = self.norm_e(e)
  
          # Chunked-cross attention if $p \in P$
          if p in self.ca_layers:
              # $H \leftarrow \text{C\small{CA}}(H, E)$
              h = self.cca[p_ca](h, e)
              # Increment chunked cross-attention index
              p_ca += 1
  
          # $H \leftarrow \text{F\small{FW}}(H)$
          h = self.ffw[p](h)
  
      # $O \leftarrow \text{R\small{EAD}}(H)$
      return self.read(h)
  
imports: import math
  from typing import Set
  import torch
  from torch import nn
  from labml.logger import inspect
label_desc: Use image node_id 2 for calling the RetroModel obj's underlying member method code with example usage: obj.forward(x, ret) and returns: self