class_docstr: A single transformer layer.
  
  Transformer layer takes input with size [s, b, h] and returns an
  output of the same size.
class_docstr_tok: ["A","single","transformer","layer",".","Transformer","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","an","output","of","the","same","size","."]
code: def forward(
      self,
      hidden_states,
      attention_mask,
      encoder_output=None,
      enc_dec_attn_mask=None,
      layer_past=None,
      get_key_value=False,
      set_inference_key_value_memory=False,
      inference_max_sequence_len=None,
      rotary_pos_emb=None,
      # list of positional embedding tensors, first one self attention, second one and third one are for cross attention (q, k)
      self_attention_relative_position_bias=None,
      cross_attention_relative_position_bias=None,
      checkpoint_core_attention=False,
  ):
      # Self attention.
      if rotary_pos_emb is not None:
          # self attention pos_emb is (q, q)
          self_attention_pos_emb = (
              rotary_pos_emb[0],
              rotary_pos_emb[0],
          )
          cross_attention_pos_emb = (
              rotary_pos_emb[1],
              rotary_pos_emb[2],
          )
      else:
          self_attention_pos_emb = None
          cross_attention_pos_emb = None
  
      if self.layer_type != LayerType.retrieval_decoder_after_self_attn:
          # hidden_states: [b, s, h]
  
          # Pre-LN: x -> LN -> MHA -> Residual -> LN -> MLP -> Residual
          # Post-LN: x -> MHA -> Residual -> LN -> MLP -> Residual -> LN
          # Normformer: x -> LN -> MHA -> LN -> Residual -> MLP (w/LN) -> Residual
  
          residual = hidden_states
          # Layer norm at the beginning of the transformer layer.
          if self.transformer_block_type in ["pre_ln", "normformer"]:
              hidden_states = self.input_layernorm(hidden_states)
  
          attention_output, attention_bias = self.self_attention(
              hidden_states,
              attention_mask,
              layer_past=layer_past,
              get_key_value=get_key_value,
              set_inference_key_value_memory=set_inference_key_value_memory,
              inference_max_sequence_len=inference_max_sequence_len,
              rotary_pos_emb=self_attention_pos_emb,
              relative_position_bias=self_attention_relative_position_bias,
              checkpoint_core_attention=checkpoint_core_attention,
          )
  
          if get_key_value:
              attention_output, presents = attention_output
  
          # If normformer, apply norm on the output of the self attention.
          if self.transformer_block_type == "normformer":
              # Normformer normalization
              attention_output = (
                  attention_output + attention_bias
                  if attention_bias is not None
                  else attention_output
              )
              attention_output = self.post_attention_normformer_norm(
                  attention_output
              )
              attention_bias = None
  
          # jit scripting for a nn.module (with dropout) is not
          # trigerring the fusion kernel. For now, we use two
          # different nn.functional routines to account for varying
          # dropout semantics during training and inference phases.
  
          if self.drop_path is None:
              bias_dropout_add_func = self._get_bias_droput_add_func(
                  transformer_block_type=self.transformer_block_type,
                  position_after="attention",
              )
              if attention_bias is not None:
                  attention_bias = attention_bias.expand_as(residual)
  
              layernorm_input = bias_dropout_add_func(
                  attention_output,
                  attention_bias,
                  residual,
                  self.hidden_dropout,
              )
          else:
              assert (
                  self.transformer_block_type != "normformer"
              ), "Normfomer doesn't support drop_path"
              out = torch.nn.functional.dropout(
                  attention_output + attention_bias,
                  p=self.hidden_dropout,
                  training=self.training,
              )
              layernorm_input = residual + self.drop_path(out)
          # print(f"Layer: {self.layer_number} Attention checksum {layernorm_input.sum()}")
  
          # Post-LN normalization after residual
          if self.transformer_block_type == "post_ln":
              normalization_output = self.input_layernorm(
                  layernorm_input
              )
              layernorm_input = normalization_output
          elif self.transformer_block_type in ["pre_ln", "normformer"]:
              # Layer norm post the self attention.
              normalization_output = self.post_attention_layernorm(
                  layernorm_input
              )
      else:
          layernorm_input, normalization_output = hidden_states
  
      if self.layer_type == LayerType.decoder_pre_mlp:
          return layernorm_input, normalization_output
  
      if (
          self.layer_type == LayerType.decoder
          or self.layer_type == LayerType.retrieval_decoder
          or self.layer_type == LayerType.retrieval_encoder
          or self.layer_type
          == LayerType.retrieval_decoder_after_self_attn
      ):
          if (
              self.layer_type == LayerType.retrieval_decoder
              or self.layer_type
              == LayerType.retrieval_decoder_after_self_attn
          ):
              attention_output, attention_bias = self.inter_attention(
                  normalization_output,
                  enc_dec_attn_mask,
                  encoder_output=encoder_output,
                  rotary_pos_emb=cross_attention_pos_emb,
                  set_inference_key_value_memory=set_inference_key_value_memory,
                  inference_max_sequence_len=inference_max_sequence_len,
                  checkpoint_core_attention=checkpoint_core_attention,
              )
          else:
              attention_output, attention_bias = self.inter_attention(
                  normalization_output,
                  enc_dec_attn_mask,
                  encoder_output=encoder_output,
                  rotary_pos_emb=cross_attention_pos_emb,
                  relative_position_bias=cross_attention_relative_position_bias,
                  checkpoint_core_attention=checkpoint_core_attention,
              )
  
          # If normformer, apply norm on the output of the self attention.
          if self.transformer_block_type == "normformer":
              # Normformer normalization
              attention_output = (
                  attention_output + attention_bias
                  if attention_bias is not None
                  else attention_output
              )
              attention_output = (
                  self.post_inter_attention_normformer_norm(
                      attention_output
                  )
              )
              attention_bias = None
  
          residual = layernorm_input
  
          bias_dropout_add_func = self._get_bias_droput_add_func(
              transformer_block_type=self.transformer_block_type,
              position_after="attention",
          )
  
          layernorm_input = bias_dropout_add_func(
              attention_output,
              attention_bias,
              residual,
              self.hidden_dropout,
          )
          # print(f"Layer: {self.layer_number} Cross-Attention checksum {layernorm_input.sum()}")
          normalization_output = self.post_inter_attention_layernorm(
              layernorm_input
          )
          # Post-LN normalization after residual
          if self.transformer_block_type == "post_ln":
              layernorm_input = normalization_output
      # MLP.
      mlp_output, mlp_bias = self.mlp(normalization_output)
  
      residual = layernorm_input
  
      if self.drop_path is None:
          bias_dropout_add_func = self._get_bias_droput_add_func(
              transformer_block_type=self.transformer_block_type,
              position_after="mlp",
          )
  
          output = bias_dropout_add_func(
              mlp_output, mlp_bias, residual, self.hidden_dropout
          )
  
      else:
          out = torch.nn.functional.dropout(
              mlp_output + mlp_bias,
              p=self.hidden_dropout,
              training=self.training,
          )
          output = residual + self.drop_path(out)
      # print(f"Layer: {self.layer_number} MLP + Dropout + Residual checksum {output.sum()}")
  
      if self.transformer_block_type == "post_ln":
          output = self.post_attention_layernorm(output)
  
      if get_key_value:
          output = [output, presents]
  
      return output
  
file_path: NeMo/nemo/collections/vision/modules/common/megatron/vision_transformer.py
imports: import torch
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.transformer import ParallelTransformer, ParallelTransformerLayer_
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
label_desc: Use image node_id 2 for calling the ParallelVisionTransformerLayer_ obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, encoder_output, enc_dec_attn_mask, layer_past, get_key_value, set_inference_key_value_memory, inference_max_sequence_len, rotary_pos_emb, _attention_relative_position_bias, cross_attention_relative_position_bias, checkpoint_core_attention) and returns: output, layernorm_input,  normalization_output