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: * `chunk_len` is the length of a chunk
  * `n_layer` is the number of layers in the encoder $L_{ ext{enc}}$
  * `ca_layers` are the layers with cross attention $P_{  ext{enc}}$
  * `d_model` is the number of features in embeddings
  * `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
docstr_tok: ["*","`","chunk_len","`","is","the","length","of","a","chunk","*","`","n_layer","`","is","the","number","of","layers","in","the","encoder","$","L_","{","ext","{","enc","}","}","$","*","`","ca_layers","`","are","the","layers","with","cross","attention","$","P_","{","ext","{","enc","}","}","$","*","`","d_model","`","is","the","number","of","features","in","embeddings","*","`","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"]
code: def __init__(
      self,
      chunk_len: int,
      n_layers: int,
      ca_layers: Set[int],
      d_model: int,
      n_heads: int,
      d_k: int,
      d_ff: int,
  ):
      """
      * `chunk_len` is the length of a chunk
      * `n_layer` is the number of layers in the encoder $L_{\text{enc}}$
      * `ca_layers` are the layers with cross attention $P_{\text{enc}}$
      * `d_model` is the number of features in embeddings
      * `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
      """
  
      super().__init__()
      self.ca_layers = ca_layers
      self.chunk_len = chunk_len
      # Cross-attention layers
      self.ca = nn.ModuleList(
          [
              CrossAttention(d_model, n_heads, d_k)
              for _ in range(len(ca_layers))
          ]
      )
      # Bi-directional self attention layers
      self.attn = nn.ModuleList(
          [
              SelfAttention(d_model, n_heads, d_k, is_causal=False)
              for _ in range(n_layers)
          ]
      )
      # Feed forward layers
      self.ffw = nn.ModuleList(
          [FeedForward(d_model, d_ff) for _ in range(n_layers)]
      )
  
      # Pre-normalization layer for $H$
      self.norm_h = 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 NearestNeighborEncoder object from inherited base classes: nn with example: obj = NearestNeighborEncoder(chunk_len, n_layers, ca_layers, d_model, n_heads, d_k, d_ff)