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: `h` are the input embeddings of shape `[batch_size, seq_len, d_model]`
  `e` are the retrieved nearest neighbors of shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
docstr_tok: ["`","h","`","are","the","input","embeddings","of","shape","`","[","batch_size",",","seq_len",",","d_model","]","`","`","e","`","are","the","retrieved","nearest","neighbors","of","shape","`","[","batch_size",",","chunks",",","neighbors",",","neighbor_len",",","d_model","]","`"]
code: def forward(self, h: torch.Tensor, e: torch.Tensor):
      """
      `h` are the input embeddings of shape `[batch_size, seq_len, d_model]`
      `e` are the retrieved nearest neighbors of shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
      """
  
      # Get shape
      batch_size, chunks, neighbors, neighbor_len, d_model = e.shape
  
      # No attention if there are no chunks (for short inputs when sampling)
      if chunks == 0:
          return h
  
      # Residual connection
      h_res = h
  
      # Remove the first `chunk_len - 1` embeddings.
      # The input pays attention to neighbors retrieved and encoded using the past tokens only;
      # so that there is no information leakage.
      # That is the retrieved neighbors from the first chunks will have information from the first chunk.
      # So by shifting the sequence to the left by `chunk_len - 1` we make sure that information only flows
      # to the right.
      h = h[:, self.chunk_len - 1 :]
      # Pre-norm
      h = self.norm(h)
      # Append empty embeddings to the end to be able to split the input into chunks
      if h.shape[1] < chunks * self.chunk_len:
          h = torch.cat(
              (
                  h,
                  h.new_zeros(
                      batch_size,
                      chunks * self.chunk_len - h.shape[1],
                      d_model,
                  ),
              ),
              dim=1,
          )
      # Reshape the input into chunks.
      h = h.reshape(batch_size, chunks, self.chunk_len, d_model)
  
      # Get query from the input
      q = self.query(h).view(*h.shape[:-1], self.n_heads, self.d_k)
      # Get keys and values from the retrieved neighbors
      k = self.key(e).view(*e.shape[:-1], self.n_heads, self.d_k)
      v = self.value(e).view(*e.shape[:-1], self.n_heads, self.d_k)
  
      # Calculate attention scores for input chunks.
      # Each chunk will pay attention to neighbors retrieved by the previous chunk.
      # This will have shape `[batch_size, chunks, heads, chunk_len, neighbors, neighbor_len]`
      attn = torch.einsum("bcihd,bcnjhd->bchinj", q, k)
      # Scale attention scores
      attn = attn * self.scale
  
      # Apply softmax over the last two dimensions `neighbors, neighbor_len`
      attn = self.softmax(attn.view(*attn.shape[:-2], -1)).view(
          attn.shape
      )
  
      # Gather values
      h = torch.einsum("bchinj,bcnjhd->bcihd", attn, v)
  
      # Change from shape `[batch_size, chunks, chunk_len, n_heads, d_k]`
      # to `[batch_size, chunks * chunk_len, n_heads * d_k]`
      h = h.reshape(batch_size, chunks * self.chunk_len, -1)
  
      # Apply final linear layer.
      # The result will have shape `[batch_size, chunks * chunk_len, d_model]`
      h = self.output(h)
  
      # Append `chunk_len - 1` zero embedding to the left; i.e. right shift it back
      h = torch.cat(
          (h.new_zeros(batch_size, self.chunk_len - 1, d_model), h),
          dim=1,
      )
  
      # Truncate and add the residual connection
      return h[:, : h_res.shape[1]] + h_res
  
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 ChunkedCrossAttention obj's underlying member method code with example usage: obj.forward(h, e) and returns: unknown, h