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","."]
docstr: This method needs to be implemented based on whether the derived class
  is "self-attn" or "cross-attn".
docstr_tok: ["This","method","needs","to","be","implemented","based","on","whether","the","derived","class","is","``","self-attn","''","or","``","cross-attn","''","."]
code: def get_query_key_value_tensors(
      self, hidden_states, key_value_states
  ):
      """
      This method needs to be implemented based on whether the derived class
      is "self-attn" or "cross-attn".
      """
  
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 5 for calling the Attention obj's underlying member method code with example usage: obj.get_query_key_value_tensors(hidden_states, key_value_states) without return types