class_docstr: ## Cross-Attention Layer $  ext{C\small{A}}$
  
  This is similar to the self-attention layer defined above, except that
  it gets keys and values from a different set of embeddings than the queries.
  
  This is used in the encoder to encode the retrieved chunks based on the
  input 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: ["#","#","Cross-Attention","Layer","$","ext","{","C\\small","{","A","}","}","$","This","is","similar","to","the","self-attention","layer","defined","above",",","except","that","it","gets","keys","and","values","from","a","different","set","of","embeddings","than","the","queries",".","This","is","used","in","the","encoder","to","encode","the","retrieved","chunks","based","on","the","input","chunks",".","*","We","do","not","use","any","explicit","positional","embeddings","here",".","We","assume","that","the","model","can","represent","positional","information","in","the","embeddings","implicitly",".","*"]
docstr: * `e` are the retrieved nearest neighbor chunk embeddings with shape
    `[batch_size, chunks, neighbors, neighbor_len, d_model]`
  * `h` are the input chunks from which the nearest neighbors were retrieved with shape
    `[batch_size, chunks, chunk_len, d_model]`. This is already normalized.
docstr_tok: ["*","`","e","`","are","the","retrieved","nearest","neighbor","chunk","embeddings","with","shape","`","[","batch_size",",","chunks",",","neighbors",",","neighbor_len",",","d_model","]","`","*","`","h","`","are","the","input","chunks","from","which","the","nearest","neighbors","were","retrieved","with","shape","`","[","batch_size",",","chunks",",","chunk_len",",","d_model","]","`",".","This","is","already","normalized","."]
code: def forward(self, e: torch.Tensor, h: torch.Tensor):
      """
      * `e` are the retrieved nearest neighbor chunk embeddings with shape
        `[batch_size, chunks, neighbors, neighbor_len, d_model]`
      * `h` are the input chunks from which the nearest neighbors were retrieved with shape
        `[batch_size, chunks, chunk_len, d_model]`. This is already normalized.
      """
  
      # Residual connection
      e_res = e
  
      # Normalize retrieved chunks
      e = self.norm(e)
  
      # Get query from the retrieved chunks
      q = self.query(e).view(*e.shape[:-1], self.n_heads, self.d_k)
      # Get keys and values from the input chunks
      k = self.key(h).view(*h.shape[:-1], self.n_heads, self.d_k)
      v = self.value(h).view(*h.shape[:-1], self.n_heads, self.d_k)
  
      # Calculate attention scores for all chunks.
      # Each retrieved neighbor will pay attention to the original chunk that retrieved it.
      # This will have shape `[batch_size, chunks, neighbors, n_heads, neighbor_len, chunk_len]`
      attn = torch.einsum("bcnihd,bcjhd->bcnhij", q, k)
      # Scale attention scores
      attn = attn * self.scale
  
      # Calculate softmax across the last dimension
      attn = self.softmax(attn)
  
      # Gather values
      e = torch.einsum("bcnhij,bcjhd->bcnihd", attn, v)
  
      # Change from shape `[batch_size, chunks, neighbors, neighbor_len, n_heads, d_k]`
      # to `[batch_size, chunks, neighbors, neighbor_len, n_heads * d_k]`
      e = e.reshape(*e.shape[:-2], -1)
  
      # Apply final linear layer.
      # The result will have shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
      e = self.output(e)
  
      # Add residual connection
      return e + e_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 CrossAttention obj's underlying member method code with example usage: obj.forward(e, h) and returns: unknown