class_docstr: Retro decoder's chunked cross attention operator.
  
  See this paper for more details: https://arxiv.org/abs/2112.04426.
  Neighboring chunks retrieved from the chunk database are used here for
  chunked-cross attention.
  
  Arguments:
    config (RetroConfig): Retro config.
  
    submodules (CrossAttentionSubmodules): Cross attention submodules.
  
    layer_number (int): Layer number within transformer block.
  
    attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
  
    encoder_block_spec (ModuleSpec): The first Retro decoder
    layer is provided with a transformer block spec to construct the
    neighbor encoder.
docstr: Cross attention for Retro decoder.
  
  Notation:
      ns : Sequence length.
      bs : Batch size.
      d  : Hidden size.
      l  : Number of chunks per sample (i.e., seq_length/chunk_length).
      m  : Number of tokens per chunk.
      k  : Number of neighbors.
      r  : Number of retrieved tokens (neighbors + continuation).
  
  Arguments:
    hidden_states (Tensor): Transformer layer hidden states.
  
    attention_mask (Tensor): Attention mask.
  
    key_value_states (Tensor): Neighbor embeddings if first decoder
    layer, else encoder output.
  
    inference_params (InferenceParams): Inference params.
docstr_tok: ["Cross","attention","for","Retro","decoder",".","Notation",":","ns",":","Sequence","length",".","bs",":","Batch","size",".","d",":","Hidden","size",".","l",":","Number","of","chunks","per","sample","(","i.e.",",","seq_length\/chunk_length",")",".","m",":","Number","of","tokens","per","chunk",".","k",":","Number","of","neighbors",".","r",":","Number","of","retrieved","tokens","(","neighbors","+","continuation",")",".","Arguments",":","hidden_states","(","Tensor",")",":","Transformer","layer","hidden","states",".","attention_mask","(","Tensor",")",":","Attention","mask",".","key_value_states","(","Tensor",")",":","Neighbor","embeddings","if","first","decoder","layer",",","else","encoder","output",".","inference_params","(","InferenceParams",")",":","Inference","params","."]
code: def forward(
      self,
      hidden_states: Tensor,
      attention_mask: Tensor,
      key_value_states: Tensor = None,
      inference_params: InferenceParams = None,
      # rotary_pos_emb: Tensor = None, # ... unsupported for retro.
  ) -> Tensor:
      """Cross attention for Retro decoder.
  
      Notation:
          ns : Sequence length.
          bs : Batch size.
          d  : Hidden size.
          l  : Number of chunks per sample (i.e., seq_length/chunk_length).
          m  : Number of tokens per chunk.
          k  : Number of neighbors.
          r  : Number of retrieved tokens (neighbors + continuation).
  
      Arguments:
        hidden_states (Tensor): Transformer layer hidden states.
  
        attention_mask (Tensor): Attention mask.
  
        key_value_states (Tensor): Neighbor embeddings if first decoder
        layer, else encoder output.
  
        inference_params (InferenceParams): Inference params.
      """
  
      # hidden_states: [ ns, bs, d ]
      # key_value_states: [ r, k*bs*l, d ]
  
      ns, bs, d = hidden_states.shape
      l = int(np.ceil(ns / self.retro_chunk_length))
  
      # Retrieve neighbors.
      if self.encoder:
          # Sequence length remainder.
          first_ns = ns % self.retro_chunk_length
  
          # Case 1: Sequence length not divisible by chunk length.
          if first_ns > 0:
              # Split sequence into first partial chunk & remaining chunks.
              first_chunk, rest_chunk = (
                  hidden_states[:first_ns],
                  hidden_states[first_ns:],
              )
  
              # Pad partial chunk with zeros.
              first_chunk = torch.nn.functional.pad(
                  first_chunk,
                  (0, 0, 0, 0, 0, self.retro_chunk_length - first_ns),
                  "constant",
                  0,
              )
  
              # Concatenate padded chunk with remaining chunks.
              chunked_output = torch.cat(
                  (first_chunk, rest_chunk), dim=0
              )  # [ l*m, bs, d ]
  
          # Case 2: Sequence length is divisible by chunk length.
          else:
              chunked_output = hidden_states  # [ l*m, bs, d ]
  
          # Chunk & permute hidden states.
          # - hidden_states:  [ l*m, bs, d ]
          # - chunked_output: [ m, bs*l, d ]
          chunked_output = (
              chunked_output.reshape(l, self.retro_chunk_length, bs, d)
              .permute(1, 2, 0, 3)
              .reshape(self.retro_chunk_length, bs * l, d)
              .contiguous()
          )
  
          # Encode neighbors. (Note: 'key_value_states' re-assigned here.)
          key_value_states = self.encoder(
              hidden_states=key_value_states,
              attention_mask=attention_mask,
              context=chunked_output,
              context_mask=None,
              inference_params=inference_params,
          )  # [ r, k*bs*l, d ]
          key_value_states = key_value_states.reshape(
              self.retro_retrieved_length * self.retro_num_neighbors,
              bs * l,
              d,
          )  # [ r*k, bs*l, d ]
  
      # Attend starting at last token of first chunk.
      pad = (ns - 1) % self.retro_chunk_length
      attending_chunks = hidden_states[pad:]
  
      # Pad attending tokens to sequence length.
      padded_chunks = torch.nn.functional.pad(
          attending_chunks,
          (0, 0, 0, 0, 0, self.retro_chunk_length - 1),
          "constant",
          0,
      )
  
      # Permute attending chunks.
      # - padded_chunks:         [ l*m, bs, d ]
      # - padded_chunked_output: [ m, bs*l, d ] (matches 'chunked_output' above)
      padded_chunked_output = padded_chunks.reshape(
          l, self.retro_chunk_length, bs, d
      ).permute(1, 2, 0, 3)
      padded_chunked_output = padded_chunked_output.reshape(
          self.retro_chunk_length, bs * l, d
      ).contiguous()
  
      # Attend to encoded neighbors.
      attention_output, attention_bias = self.attn(
          padded_chunked_output,
          None,
          key_value_states=key_value_states,
      )
  
      # Return dimensions for bias-dropout step.
      return {
          "ns": ns,
          "bs": bs,
          "d": d,
          "l": l,
          "pad": pad,
          "attention_output": attention_output,  # [ m, bs*l, d ]
          "attention_bias": attention_bias,  # [ d ]
          "context": key_value_states,  # [ r*k, bs*l, d ]
      }
  
imports: from functools import partial
  from typing import Callable
  import numpy
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams
  from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
  from megatron.core.models.retro.base_attention import BaseRetroCrossAttention
  from megatron.core.models.retro.config import RetroConfig
  from megatron.core.transformer import ModuleSpec
  from megatron.core.transformer.attention import CrossAttentionSubmodules
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_block import TransformerBlock
label_desc: Use image node_id 2 for calling the RetroDecoderCrossAttention obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, key_value_states, inference_params) and returns: dict