class_docstr: Attention layer abstract class.
  
  This layer only contains common modules required for the "self attn" and
  "cross attn" specializations.
class_docstr_tok: ["Attention","layer","abstract","class",".","This","layer","only","contains","common","modules","required","for","the","``","self","attn","''","and","''","cross","attn","''","specializations","."]
code: def forward(
      self,
      hidden_states,
      attention_mask,
      key_value_states=None,
      inference_params=None,
      rotary_pos_emb=None,
  ):
      # hidden_states: [sq, b, h]
  
      # For self attention we just duplicate the rotary_pos_emb if it isn't already
      if rotary_pos_emb is not None and not isinstance(
          rotary_pos_emb, tuple
      ):
          rotary_pos_emb = (rotary_pos_emb,) * 2
  
      # =====================
      # Query, Key, and Value
      # =====================
      # Get the query, key and value tensors based on the type of attention -
      # self or cross attn.
      query, key, value = self.get_query_key_value_tensors(
          hidden_states, key_value_states
      )
  
      # ===================================================
      # Adjust key, value, and rotary_pos_emb for inference
      # ===================================================
      (
          key,
          value,
          rotary_pos_emb,
          attn_mask_type,
      ) = self._adjust_key_value_for_inference(
          inference_params, key, value, rotary_pos_emb
      )
  
      # ================================================
      # relative positional embedding (rotary embedding)
      # ================================================
      if rotary_pos_emb is not None:
          q_pos_emb, k_pos_emb = rotary_pos_emb
          query = apply_rotary_pos_emb(query, q_pos_emb)
          key = apply_rotary_pos_emb(key, k_pos_emb)
          # TODO, can apply positional embedding to value_layer so it has
          # absolute positional embedding.
          # otherwise, only relative positional embedding takes effect
          # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb)
  
      # ==================================
      # core attention computation
      # ==================================
  
      if self.checkpoint_core_attention:
          core_attn_out = self._checkpointed_attention_forward(
              query,
              key,
              value,
              attention_mask,
              attn_mask_type=attn_mask_type,
          )
      else:
          core_attn_out = self.core_attention(
              query,
              key,
              value,
              attention_mask,
              attn_mask_type=attn_mask_type,
          )
  
      # =================
      # Output. [sq, b, h]
      # =================
  
      output, bias = self.linear_proj(core_attn_out)
  
      return output, bias
  
imports: from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb
  from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
  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 divide
  from .enums import AttnMaskType
  from .transformer_config import TransformerConfig
  from .utils import make_sharded_tensors_for_checkpoint
label_desc: Use image node_id 6 for calling the Attention obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, key_value_states, inference_params, rotary_pos_emb) and returns: output,  bias