class_docstr: Transformer encoder model.
      
code: def forward(
      self,
      enc_input,
      enc_attn_mask,
      context_attn_mask=None,
      encoder_output=None,
      layer_past=None,
      get_key_value=False,
      set_inference_key_value_memory=False,  # when doing inference, set this to true to allocate all the cached matrix. later set false to do incremental inference
      inference_max_sequence_len=None,
      neighbors=2,
  ):
      # expected enc_input shape [batch, num_chunks, num_neighbors, retrieval_seq_len, dim]
      # expected enc_attn_mask shape [batch, num_chunks, num_neighbors, retrieval_seq_len]
      # expected encoder_output shape [batch, seq_len, dim]
  
      # batch, seq_len, dim
      b, n, dim = encoder_output.shape
  
      if set_inference_key_value_memory:
          # run once to setup the cache
          chunk_start = 0
          num_seq_chunks = n // self.chunk_size
          num_chunks = inference_max_sequence_len // self.chunk_size
          self.cache_output = self._allocate_memory(
              b,
              num_chunks,
              neighbors,
              self.chunk_size * 2,
              dim,
              dtype=encoder_output.dtype,
          )
          self.seq_pos_in_chunk = n
          self.current_chunk = n // self.chunk_size
          self.encoder_output = self._allocate_memory(
              b, self.chunk_size, dim, dtype=encoder_output.dtype
          )
          self.context_attn_mask = self._allocate_memory(
              b, self.chunk_size, dtype=context_attn_mask.dtype
          )
          self.context_attn_mask
          chunk_beg = self.chunk_size * num_seq_chunks
          chunk_end = (
              self.chunk_size * num_seq_chunks
              + self.seq_pos_in_chunk % self.chunk_size
          )
          # store the remainders
          self.encoder_output[
              :, : self.seq_pos_in_chunk % self.chunk_size, :
          ] = encoder_output[:, chunk_beg:chunk_end, :]
          self.context_attn_mask[
              :, : self.seq_pos_in_chunk % self.chunk_size
          ] = context_attn_mask[:, chunk_beg:chunk_end]
      elif inference_max_sequence_len is not None:
          # second time of running
          # only support one token at a time
          assert n == 1
          self.seq_pos_in_chunk += n
          self.current_chunk = self.seq_pos_in_chunk // self.chunk_size
          # if exceed the chunk size
          pos_beg = (self.seq_pos_in_chunk - 1) % self.chunk_size
          # if self.seq_pos_in_chunk - 1 >= self.chunk_size:
          #     self.current_chunk += 1
          #     self.seq_pos_in_chunk -= self.chunk_size
          chunk_start = self.current_chunk - 1
          self.encoder_output[
              :, pos_beg : pos_beg + 1, :
          ] = encoder_output
          self.context_attn_mask[
              :, pos_beg : pos_beg + 1
          ] = context_attn_mask[
              :, self.seq_pos_in_chunk - 1 : self.seq_pos_in_chunk
          ]
          encoder_output = self.encoder_output[:, : pos_beg + 1, :]
          context_attn_mask = self.context_attn_mask[:, : pos_beg + 1]
          num_seq_chunks = 1
          if not self.seq_pos_in_chunk % self.chunk_size == 0:
              # still accumulate the encoder_output
              # return the cached results
              if self.current_chunk == 0:
                  return None
              return self.cache_output[:, : self.current_chunk]
          if enc_input is not None:
              # only need one chunk for the later calculation
              enc_input = enc_input[
                  :, self.current_chunk - 1 : self.current_chunk
              ]
              enc_attn_mask = enc_attn_mask[
                  :, self.current_chunk - 1 : self.current_chunk
              ]
  
      if enc_input is None:
          return None
  
      _, k, r, rn, _ = enc_input.shape
  
      assert r == neighbors
      if inference_max_sequence_len is None:
          num_seq_chunks = n // self.chunk_size
          assert (
              k == num_seq_chunks
          ), f"sequence requires {num_seq_chunks} retrieved chunks, but only {k} passed in"
      else:
          pass
  
      seq_index = num_seq_chunks * self.chunk_size
  
      retrieved = rearrange(enc_input, "b k r n d -> n (b k r) d")
      enc_attn_mask = rearrange(enc_attn_mask, "b k r n -> (b k r) n")
      # embed_as_context = repeat(encoder_output[:, :seq_index], 'b (k n) d -> (b k r) n d', n=self.chunk_size, r=r)
      # context_attn_mask = repeat(context_attn_mask[:, :seq_index], 'b (k n) -> (b k r) n', n=self.chunk_size, r=r)
  
      if (
          inference_max_sequence_len is not None
          and not set_inference_key_value_memory
      ):
          embed_as_context = repeat(
              encoder_output[:, :seq_index],
              "b (k n) d -> n (b k r) d",
              n=pos_beg + 1,
              r=r,
          )
          context_attn_mask = repeat(
              context_attn_mask[:, :seq_index],
              "b (k n) -> (b k r) n",
              n=pos_beg + 1,
              r=r,
          )
      else:
          embed_as_context = repeat(
              encoder_output[:, :seq_index],
              "b (k n) d -> n (b k r) d",
              n=self.chunk_size,
              r=r,
          )
          context_attn_mask = repeat(
              context_attn_mask[:, :seq_index],
              "b (k n) -> (b k r) n",
              n=self.chunk_size,
              r=r,
          )
  
      if not self.turn_off_rop:
          if (
              inference_max_sequence_len is not None
              and not set_inference_key_value_memory
          ):
              cross_attn_k_pos_emb = self.rotary_pos_emb(
                  n % self.chunk_size, offset=pos_beg
              )
          else:
              cross_attn_k_pos_emb = self.rotary_pos_emb(
                  self.chunk_size, offset=0
              )
          cross_attn_q_pos_emb = self.rotary_pos_emb(rn, offset=0)
          attn_pos_emb = (
              cross_attn_q_pos_emb,
              cross_attn_q_pos_emb,
              cross_attn_k_pos_emb,
          )
      else:
          attn_pos_emb = None
  
      # # convert to Megatron mask
      enc_attn_mask_3d = build_attention_mask_3d(
          source_mask=enc_attn_mask,
          target_mask=enc_attn_mask,
          attn_mask_type=AttnMaskType.padding,
      )
      enc_attn_mask_3d = enc_attn_mask_3d[:, None, :, :]
  
      enc_dec_attn_mask_3d = build_attention_mask_3d(
          source_mask=enc_attn_mask,
          target_mask=context_attn_mask,
          attn_mask_type=AttnMaskType.padding,
      )
      enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
  
      # transformer encoder
      enc_output = self.model(
          retrieved,
          enc_attn_mask_3d,
          layer_past=layer_past,
          get_key_value=get_key_value,
          encoder_output=embed_as_context,
          enc_dec_attn_mask=enc_dec_attn_mask_3d,
          rotary_pos_emb=attn_pos_emb,
      )
      # revert back to original retrieved shape
      enc_output = rearrange(
          enc_output, "n (b k r) d -> b k r n d", b=b, k=k
      )
  
      if inference_max_sequence_len is not None:
          # update encoded for current chunk
          self.cache_output[
              :, chunk_start : self.current_chunk, :, :, :
          ] = enc_output
          # read all encodings
          enc_output = self.cache_output[:, : self.current_chunk]
      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 MegatronRetrievalTransformerEncoderModule obj's underlying member method code with example usage: obj.forward(enc_input, enc_attn_mask, context_attn_mask, encoder_output, layer_past, get_key_value, set_inference_key_value_memory, inference_max_sequence_len, neighbors) and returns: enc_output, None, self, None