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,
      context=None,
      context_mask=None,
      rotary_pos_emb=None,
      inference_params=None,
  ):
      # hidden_states: [s, b, h]
  
      # Residual connection.
      residual = hidden_states
  
      # Optional Input Layer norm
      input_layernorm_output = self.input_layernorm(hidden_states)
  
      # Self attention.
      attention_output_with_bias = self.self_attention(
          input_layernorm_output,
          attention_mask=attention_mask,
          inference_params=inference_params,
          rotary_pos_emb=rotary_pos_emb,
      )
  
      # TODO: could we move `bias_dropout_add_exec_handler` itself
      # inside the module provided in the `bias_dropout_add_spec` module?
      with self.bias_dropout_add_exec_handler():
          hidden_states = self.self_attn_bda(
              self.training, self.config.bias_dropout_fusion
          )(attention_output_with_bias, residual, self.hidden_dropout)
  
      # Residual connection.
      residual = hidden_states
  
      # Optional Layer norm after self-attention
      pre_cross_attn_layernorm_output = self.pre_cross_attn_layernorm(
          hidden_states
      )
  
      # Cross attention.
      attention_output_with_bias = self.cross_attention(
          pre_cross_attn_layernorm_output,
          attention_mask=context_mask,
          key_value_states=context,
          inference_params=inference_params,
      )
  
      if (
          isinstance(attention_output_with_bias, dict)
          and "context" in attention_output_with_bias
      ):
          context = attention_output_with_bias["context"]
  
      # TODO: could we move `bias_dropout_add_exec_handler` itself
      # inside the module provided in the `bias_dropout_add_spec` module?
      with self.bias_dropout_add_exec_handler():
          hidden_states = self.cross_attn_bda(
              self.training, self.config.bias_dropout_fusion
          )(attention_output_with_bias, residual, self.hidden_dropout)
  
      # Residual connection.
      residual = hidden_states
  
      # Optional Layer norm post the cross-attention.
      pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
  
      # MLP.
      mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
  
      # TODO: could we move `bias_dropout_add_exec_handler` itself
      # inside the module provided in the `bias_dropout_add_spec` module?
      with self.bias_dropout_add_exec_handler():
          hidden_states = self.mlp_bda(
              self.training, self.config.bias_dropout_fusion
          )(mlp_output_with_bias, residual, self.hidden_dropout)
  
      # Jit compiled function creates 'view' tensor. This tensor
      # potentially gets saved in the MPU checkpoint function context,
      # which rejects view tensors. While making a viewless tensor here
      # won't result in memory savings (like the data loader, or
      # p2p_communication), it serves to document the origin of this
      # 'view' tensor.
      output = make_viewless_tensor(
          inp=hidden_states,
          requires_grad=hidden_states.requires_grad,
          keep_graph=True,
      )
  
      return output, context
  
file_path: megatron-lm/megatron/core/transformer/transformer_layer.py
imports: from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import make_viewless_tensor
label_desc: Use image node_id 3 for calling the TransformerLayer obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, context, context_mask, rotary_pos_emb, inference_params) and returns: output,  context