class_docstr: Transformer decoder model.
      
code: def forward(
      self,
      dec_input,
      dec_attn_mask,
      retrieved_attn_mask=None,
      retrieved_emb=None,
      layer_past=None,
      get_key_value=False,
      eod_positions=None,  # this is a tuple of eod positions returned from tensor.where(tensor == eod_id)
      set_inference_key_value_memory=False,
      inference_max_sequence_len=None,
  ):
      # expected dec_input shape [batch, seq_len, dim]
      # expected dec_attn_mask shape [batch, seq_len]
      # expected retrieved_input shape [batch, num_chunks, num_neighbors, retrival_seq_len, dim]
      # expected retrieved_attn_mask shape [batch, num_chunks, num_neighbors, retrival_seq_len]
  
      # batch, seq_len, dim
      if isinstance(dec_input, tuple):
          n, _, _ = dec_input[1].shape
      else:
          _, n, _ = dec_input.shape
  
      if set_inference_key_value_memory:
          # seq_index = (n // chunk_size) * chunk_size
          self.current_len = n
          num_seq_chunks = self.current_len // self.chunk_size
      elif inference_max_sequence_len is not None:
          # only handles single token increment
          assert n == 1
          self.current_len += n
          num_seq_chunks = self.current_len // self.chunk_size
      else:
          # this is normal forward without inference
          num_seq_chunks = n // self.chunk_size
  
      if retrieved_emb is not None:
          b, k, r, rn, dim = retrieved_emb.shape
          assert (
              k == num_seq_chunks
          ), f"sequence requires {num_seq_chunks} retrieved chunks, but only {k} passed in"  # need to add extra chunk size, since it will be shifted
  
      if not self.turn_off_rop:
          if set_inference_key_value_memory:
              self_attn_emb = self.rotary_pos_emb(self.current_len)
          elif inference_max_sequence_len is not None:
              self_attn_emb = self.rotary_pos_emb(self.current_len)
          else:
              self_attn_emb = self.rotary_pos_emb(n)
          if retrieved_emb is not None:
              # -63, -62, ... 63  will be cut into -> [0, ... 63] in the chunk cross attention layer
              cross_attn_q_pos_emb = self.rotary_pos_emb(
                  self.chunk_size * 2 - 1, offset=-self.chunk_size + 1
              )
              if self.version == 1:
                  cross_attn_k_pos_emb = self.rotary_pos_emb(
                      rn, offset=0
                  )
              elif self.version > 1:
                  # the first 64 tokens in retrieved is from the last chunk, align the continuation part with the query tokens
                  # use the following in the future. [-63, -62, ..., 63, 64]
                  cross_attn_k_pos_emb = self.rotary_pos_emb(
                      rn, offset=-self.chunk_size + 1
                  )
              else:
                  raise ValueError(
                      f"incorrect version number {self.version}"
                  )
              attn_pos_emb = (
                  self_attn_emb,
                  cross_attn_q_pos_emb,
                  cross_attn_k_pos_emb,
              )
          else:
              attn_pos_emb = (self_attn_emb, None, None)
      else:
          attn_pos_emb = None
  
      dec_attn_mask_3d = self._calculate_dec_att_mask(
          dec_attn_mask, eod_positions
      )
  
      if retrieved_emb is not None:
          # need to shift the dec_attn_mask as first causal_padding elements are ignored
          # also pad it to be the multiple of self.chunk_size
          causal_padding = self.chunk_size - 1
          reminder = (
              self.chunk_size - (dec_attn_mask.shape[1] + 1)
          ) % self.chunk_size
          dec_attn_mask = F.pad(
              dec_attn_mask, (-causal_padding, reminder), value=False
          )
  
          dec_attn_mask = rearrange(
              dec_attn_mask, "b (k n) -> (b k) n", k=k
          )
          retrieved_attn_mask = rearrange(
              retrieved_attn_mask, "b k r n -> (b k) (r n)"
          )
  
          enc_dec_attn_mask_3d = build_attention_mask_3d(
              source_mask=dec_attn_mask,
              target_mask=retrieved_attn_mask,
              attn_mask_type=AttnMaskType.padding,
          )
          enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
      else:
          enc_dec_attn_mask_3d = None
  
      # transformer encoder
      if not isinstance(dec_input, tuple):
          dec_input = rearrange(
              dec_input, "b s d -> s b d"
          ).contiguous()
      enc_output = self.model(
          dec_input,
          dec_attn_mask_3d,
          layer_past=layer_past,
          get_key_value=get_key_value,
          encoder_output=None,
          retrieved_emb=retrieved_emb,
          enc_dec_attn_mask=enc_dec_attn_mask_3d,
          rotary_pos_emb=attn_pos_emb,
          set_inference_key_value_memory=set_inference_key_value_memory,
          inference_max_sequence_len=inference_max_sequence_len,
      )
      # enc_output = rearrange(dec_input, 's b d -> b s d')
      return enc_output
  
imports: import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import RotaryEmbedding
  from nemo.collections.nlp.modules.common.megatron.transformer import ParallelTransformer
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, build_attention_mask_3d
label_desc: Use image node_id 4 for calling the MegatronRetrievalTransformerDecoderModule obj's underlying member method code with example usage: obj.forward(dec_input, dec_attn_mask, retrieved_attn_mask, retrieved_emb, layer_past, get_key_value, eod_positions, set_inference_key_value_memory, inference_max_sequence_len) and returns: enc_output