class_docstr: ## Chunked Cross-Attention Layer $  ext{C\small{CA}}$
  
  This is similar to the cross-attention layer defined above.
  
  This is used in the decoder to pay attention to the retrieved neighbor chunks.
  
  *We do not use any explicit positional embeddings here.
  We assume that the model can represent positional information in the embeddings implicitly.*
class_docstr_tok: ["#","#","Chunked","Cross-Attention","Layer","$","ext","{","C\\small","{","CA","}","}","$","This","is","similar","to","the","cross-attention","layer","defined","above",".","This","is","used","in","the","decoder","to","pay","attention","to","the","retrieved","neighbor","chunks",".","*","We","do","not","use","any","explicit","positional","embeddings","here",".","We","assume","that","the","model","can","represent","positional","information","in","the","embeddings","implicitly",".","*"]
docstr: * `d_model` is the number of features in transformer embeddings
  * `n_heads` is the number of attention heads
  * `d_k` is the number of features per head
  * `chunk_len` is the length of a chunk
docstr_tok: ["*","`","d_model","`","is","the","number","of","features","in","transformer","embeddings","*","`","n_heads","`","is","the","number","of","attention","heads","*","`","d_k","`","is","the","number","of","features","per","head","*","`","chunk_len","`","is","the","length","of","a","chunk"]
code: def __init__(
      self, d_model: int, n_heads: int, d_k: int, chunk_len: int
  ):
      """
      * `d_model` is the number of features in transformer embeddings
      * `n_heads` is the number of attention heads
      * `d_k` is the number of features per head
      * `chunk_len` is the length of a chunk
      """
  
      super().__init__()
  
      self.chunk_len = chunk_len
      self.n_heads = n_heads
      self.d_k = d_k
  
      # To scale attentions before softmax by $\frac{1}{\sqrt{d_k}}$
      self.scale = 1 / math.sqrt(self.d_k)
  
      # Linear layers for query, key and value heads.
      self.query = nn.Linear(d_model, n_heads * d_k)
      self.key = nn.Linear(d_model, n_heads * d_k)
      self.value = nn.Linear(d_model, n_heads * d_k)
  
      # Pre-norm layer for the query embeddings. The paper uses RMSNorm instead.
      self.norm = nn.LayerNorm(d_model)
  
      # Softmax for attention probabilities
      self.softmax = nn.Softmax(dim=-1)
  
      # Final linear layer
      self.output = nn.Linear(n_heads * d_k, 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 ChunkedCrossAttention object from inherited base classes: nn with example: obj = ChunkedCrossAttention(d_model, n_heads, d_k, chunk_len)