class_docstr: ## Nearest Neighbor Encoder $       ext{E\small{NCODER}}(   ext{R\small{ET}}(C_u)_{1 \le u \le l}, H)$
  
  This module encodes the retrieved nearest neighbors
class_docstr_tok: ["#","#","Nearest","Neighbor","Encoder","$","ext","{","E\\small","{","NCODER","}","}","(","ext","{","R\\small","{","ET","}","}","(","C_u",")","_","{","1","\\le","u","\\le","l","}",",","H",")","$","This","module","encodes","the","retrieved","nearest","neighbors"]
docstr: * `e` are token embeddings of the retrieved nearest neighbors,
   $      ext{E\small{MB}}ig(    ext{R\small{ET}}(C_u)_{1 \le u \le l}ig)$
   of shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
  
  * `h` is are the input token embeddings, $H$
   of shape `[batch_size, seq_len, d_model]`
  
  *The chunks $u \in [1, l]$ and neighbors $j \in [1, k]$ are processed in parallel.*
docstr_tok: ["*","`","e","`","are","token","embeddings","of","the","retrieved","nearest","neighbors",",","$","ext","{","E\\small","{","MB","}","}","\big","(","ext","{","R\\small","{","ET","}","}","(","C_u",")","_","{","1","\\le","u","\\le","l","}","\big",")","$","of","shape","`","[","batch_size",",","chunks",",","neighbors",",","neighbor_len",",","d_model","]","`","*","`","h","`","is","are","the","input","token","embeddings",",","$","H","$","of","shape","`","[","batch_size",",","seq_len",",","d_model","]","`","*","The","chunks","$","u","\\in","[","1",",","l","]","$","and","neighbors","$","j","\\in","[","1",",","k","]","$","are","processed","in","parallel",".","*"]
code: def forward(self, e: torch.Tensor, h: torch.Tensor):
      """
      * `e` are token embeddings of the retrieved nearest neighbors,
       $\text{E\small{MB}}\big(\text{R\small{ET}}(C_u)_{1 \le u \le l}\big)$
       of shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
  
      * `h` is are the input token embeddings, $H$
       of shape `[batch_size, seq_len, d_model]`
  
      *The chunks $u \in [1, l]$ and neighbors $j \in [1, k]$ are processed in parallel.*
      """
  
      # Get shape
      batch_size, chunks, neighbors, neighbor_len, d_model = e.shape
  
      # $(H_u)_{u \in [1, l]} \leftarrow \text{S\small{PLIT}}(H)$
      h_split = h[:, : self.chunk_len * chunks, :].reshape(
          batch_size, chunks, self.chunk_len, d_model
      )
  
      # Pre-norm
      h_split = self.norm_h(h_split)
  
      # Keep the index of the cross attention layer
      p_ca = 0
      # For all layers $p' \in [1, L_{\text{enc}}]$
      for p in range(len(self.attn)):
          # Bi-directional self attention
          # $E^j_u \leftarrow \text{A\small{TTN}}_{\text{enc}}(E^j_u)$
          e = self.attn[p](e.view(-1, neighbor_len, d_model)).view(
              e.shape
          )
  
          # Cross attention if $p' \in P_{\text{enc}}$
          if p in self.ca_layers:
              # $E^j_u \leftarrow \text{C\small{A}}_{\text{enc}}(E^j_u, H_u)$
              e = self.ca[p_ca](e, h_split)
              # Incremnt the cross attention index
              p_ca += 1
  
          # Feed forward layer $E^j_u \leftarrow \text{F\small{FW}}_{\text{enc}}(E^j_u)$
          e = self.ffw[p](e)
  
      # return $E$
      return e
  
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 NearestNeighborEncoder obj's underlying member method code with example usage: obj.forward(e, h) and returns: e