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 __init__(
      self,
      config: TransformerConfig,
      submodules: Union[
          SelfAttentionSubmodules, CrossAttentionSubmodules
      ],
      layer_number: int,
      attn_mask_type: AttnMaskType,
      attention_type: str,
  ):
      super().__init__(config=config)
  
      self.config = config
      self.layer_number = layer_number
      self.attn_mask_type = attn_mask_type
      self.attention_type = attention_type
  
      # For normal attention without groups, num_query_groups == num_attention_heads,
      # so these two will be the same
      self.query_projection_size = (
          self.config.kv_channels * self.config.num_attention_heads
      )
      self.kv_projection_size = (
          self.config.kv_channels * self.config.num_query_groups
      )
  
      # Per attention head and per partition values.
      world_size = parallel_state.get_tensor_model_parallel_world_size()
      self.hidden_size_per_attention_head = divide(
          self.query_projection_size, self.config.num_attention_heads
      )
      self.num_attention_heads_per_partition = divide(
          self.config.num_attention_heads, world_size
      )
      self.num_query_groups_per_partition = divide(
          self.config.num_query_groups, world_size
      )
  
      self.core_attention = build_module(
          submodules.core_attention,
          config=self.config,
          layer_number=self.layer_number,
          attn_mask_type=self.attn_mask_type,
          attention_type=self.attention_type,
      )
  
      self.checkpoint_core_attention = (
          self.config.recompute_granularity == "selective"
      )
  
      # Output.
      self.linear_proj = build_module(
          submodules.linear_proj,
          self.query_projection_size,
          self.config.hidden_size,
          config=self.config,
          init_method=self.config.output_layer_init_method,
          bias=self.config.add_bias_linear,
          input_is_parallel=True,
          skip_bias_add=True,
          is_expert=False,
          tp_comm_buffer_name="proj",
      )
  
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 1 to create a new Attention object from inherited base classes: MegatronModule, ABC with example: obj = Attention(config, submodules, layer_number, attn_mask_type, attention_type)