class_docstr: Parallel self-attention layer abstract class.
  
  Self-attention layer takes input with size [s, b, h]
  and returns output of the same size.
class_docstr_tok: ["Parallel","self-attention","layer","abstract","class",".","Self-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
code: def __init__(
      self,
      config: ModelParallelConfig,
      init_method,
      output_layer_init_method,
      layer_number,
      num_attention_heads,
      hidden_size,
      attention_type=AttnType.self_attn,
      attn_mask_type=AttnMaskType.padding,
      precision=16,
      apply_query_key_layer_scaling=False,
      kv_channels=None,
      masked_softmax_fusion=True,
      attention_dropout=0.1,
      layer_type=None,
      megatron_legacy=False,
      bias=True,
      headscale=False,
      position_embedding_type="learned_absolute",
      multi_query_attention=False,
      normalize_attention_scores=True,
      use_flash_attention=False,
  ):
      super(ParallelAttention, self).__init__(config=config)
      self.layer_number = max(1, layer_number)
      self.attention_type = attention_type
      self.attn_mask_type = attn_mask_type
      self.normalize_attention_scores = normalize_attention_scores
      self.position_embedding_type = position_embedding_type
      self.multi_query_attention = multi_query_attention
      self.use_flash_attention = use_flash_attention
  
      self.megatron_legacy = megatron_legacy
  
      self.set_accepted_adapter_types(
          [
              InfusedAdapterConfig._target_,
              LoraKQVAdapterConfig._target_,
              LoraQAdapterConfig._target_,
              LoraKVAdapterConfig._target_,
              LoraKQVAdapterWeightTyingConfig._target_,
          ]
      )
  
      if kv_channels is None:
          assert (
              hidden_size % num_attention_heads == 0
          ), "hidden_size must be divisible by num_attention_heads if kv_channels is None"
          kv_channels = hidden_size // num_attention_heads
      projection_size = kv_channels * num_attention_heads
  
      # Per attention head and per partition values.
      world_size = parallel_state.get_tensor_model_parallel_world_size()
      self.hidden_size_per_attention_head = safe_divide(
          projection_size, num_attention_heads
      )
      self.num_attention_heads_per_partition = safe_divide(
          num_attention_heads, world_size
      )
      self.num_attention_heads_partition_offset = (
          self.num_attention_heads_per_partition
          * parallel_state.get_tensor_model_parallel_rank()
      )
  
      # Strided linear layer.
      if attention_type == AttnType.self_attn:
          self.query_key_value = tensor_parallel.ColumnParallelLinear(
              hidden_size,
              3 * projection_size,
              config=config,
              gather_output=False,
              init_method=init_method,
              bias=bias,
          )
      else:
          assert attention_type == AttnType.cross_attn
          self.query = tensor_parallel.ColumnParallelLinear(
              hidden_size,
              projection_size,
              config=config,
              gather_output=False,
              init_method=init_method,
              bias=bias,
          )
  
          self.key_value = tensor_parallel.ColumnParallelLinear(
              hidden_size,
              2 * projection_size,
              config=config,
              gather_output=False,
              init_method=init_method,
              bias=bias,
          )
  
      self.core_attention = CoreAttention(
          config=config,
          layer_number=self.layer_number,
          num_attention_heads=num_attention_heads,
          hidden_size=hidden_size,
          attention_type=self.attention_type,
          attn_mask_type=self.attn_mask_type,
          precision=precision,
          apply_query_key_layer_scaling=apply_query_key_layer_scaling,
          kv_channels=kv_channels,
          masked_softmax_fusion=masked_softmax_fusion,
          attention_dropout=attention_dropout,
          multi_query_attention=multi_query_attention,
          normalize_attention_scores=normalize_attention_scores,
          position_embedding_type=position_embedding_type,
          use_flash_attention=use_flash_attention,
      )
  
      # Output.
      self.dense = tensor_parallel.RowParallelLinear(
          projection_size,
          hidden_size,
          config=config,
          input_is_parallel=True,
          init_method=output_layer_init_method,
          skip_bias_add=True,
          bias=bias,
      )
  
      self.headscale = headscale
      if headscale:
          self.head_scale_tensor = torch.nn.Parameter(
              torch.ones(
                  1, self.num_attention_heads_per_partition, 1, 1
              ),
              requires_grad=True,
          )
  
      # Inference key-value memory
      self.inference_key_memory = None
      self.inference_value_memory = None
      self.inference_current_sequence_len = 0
  
      # relative position embedding
      self.layer_type = layer_type
  
imports: import math
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, LoraKQVAdapterWeightTyingConfig, LoraKVAdapterConfig, LoraQAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_softmax import MatchedScaleMaskSoftmax
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import XPOSPositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.position_embedding.rotary_position_embedding import apply_rotary_pos_emb
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, _cast_if_autocast_enabled, attention_mask_func
  from nemo.core import adapter_mixins
label_desc: Use image node_id 1 to create a new ParallelAttention object from inherited base classes: MegatronModule, adapter_mixins with example: obj = ParallelAttention(config, init_method, output_layer_init_method, layer_number, num_attention_heads, hidden_size, attention_type, attn_mask_type, precision, apply_query_key_layer_scaling, kv_channels, masked_softmax_fusion, attention_dropout, layer_type, megatron_legacy, bias, headscale, position_embedding_type, multi_query_attention, normalize_attention_scores, use_flash_attention)