class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
code: def forward(
      self,
      hidden_states: Tensor,
      attention_mask: Tensor,
      context: Tensor = None,
      context_mask: Tensor = None,
      rotary_pos_emb: Tensor = None,
      inference_params: InferenceParams = None,
  ):
      # hidden_states (float): [s, b, h]
      # attention_mask (bool): [1, 1, s, s]
  
      if not self.pre_process:
          # See set_input_tensor()
          hidden_states = self.input_tensor
  
      # Viewless tensor.
      # - We only need to create a viewless tensor in the case of micro batch
      #   size (mbs) == 1, since in this case, 'hidden_states.transpose()'
      #   above creates a view tensor, and '.contiguous()' is a pass-through.
      #   For mbs >= 2, '.contiguous()' creates a new tensor, eliminating
      #   the need to make it viewless.
      #
      #   However, we don't explicitly check mbs == 1 here because
      #   make_viewless_tensor() has negligible overhead when its input
      #   is already viewless.
      #
      # - For the 'else' case above, calling make_viewless_tensor() here is
      #   likely redundant, since p2p_communication.py (likely originator)
      #   already creates viewless tensors. That said, make_viewless_tensor()
      #   is called here to be future-proof and corner-case-proof.
      hidden_states = make_viewless_tensor(
          inp=hidden_states,
          requires_grad=True,
          keep_graph=True,
      )
  
      if self.config.sequence_parallel:
          rng_context = tensor_parallel.get_cuda_rng_tracker().fork()
      else:
          rng_context = nullcontext()
  
      if self.config.fp8:
          import transformer_engine  # To keep out TE dependency when not training in fp8
  
          if self.config.fp8 == "e4m3":
              fp8_format = transformer_engine.common.recipe.Format.E4M3
          elif self.config.fp8 == "hybrid":
              fp8_format = (
                  transformer_engine.common.recipe.Format.HYBRID
              )
          else:
              raise ValueError(
                  "E4M3 and HYBRID are the only supported FP8 formats."
              )
  
          fp8_recipe = transformer_engine.common.recipe.DelayedScaling(
              margin=self.config.fp8_margin,
              interval=self.config.fp8_interval,
              fp8_format=fp8_format,
              amax_compute_algo=self.config.fp8_amax_compute_algo,
              amax_history_len=self.config.fp8_amax_history_len,
              override_linear_precision=(
                  False,
                  False,
                  not self.config.fp8_wgrad,
              ),
          )
          fp8_group = None
          if parallel_state.model_parallel_is_initialized():
              fp8_group = parallel_state.get_amax_reduction_group(
                  with_context_parallel=True
              )
          fp8_context = transformer_engine.pytorch.fp8_autocast(
              enabled=True, fp8_recipe=fp8_recipe, fp8_group=fp8_group
          )
      else:
          fp8_context = nullcontext()
  
      with rng_context and fp8_context:
          # Forward pass.
          if self.config.recompute_granularity == "full":
              hidden_states = self._checkpointed_forward(
                  hidden_states=hidden_states,
                  attention_mask=attention_mask,
                  context=context,
                  context_mask=context_mask,
                  rotary_pos_emb=rotary_pos_emb,
              )
          else:
              for layer in self.layers:
                  hidden_states, context = layer(
                      hidden_states=hidden_states,
                      attention_mask=attention_mask,
                      context=context,
                      context_mask=context_mask,
                      rotary_pos_emb=rotary_pos_emb,
                      inference_params=inference_params,
                  )
  
      # Final layer norm.
      if self.post_process and self.post_layer_norm:
          hidden_states = self.final_layernorm(hidden_states)
  
      return hidden_states
  
imports: import re
  from contextlib import nullcontext
  from dataclasses import dataclass
  from typing import List, Union
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams, parallel_state, tensor_parallel
  from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
  from megatron.core.transformer.custom_layers.transformer_engine import TENorm
  from megatron.core.transformer.enums import AttnMaskType
  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.transformer.transformer_layer import TransformerLayer
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_viewless_tensor
label_desc: Use image node_id 6 for calling the TransformerBlock 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: hidden_states