class_docstr: Parallel self-attention layer abstract class.
  Self-attention layer takes input with size [b, s, h]
  and returns output of the same size.
  
  Args:
      hidden_size (int): hidden size.
      num_attention_heads (int): number of attention heads.
      attention_dropout (float): dropout probability for attention layer.
      attention_mask_func (:class:`typing.Callable`): Mask function to be applied.
      layer_number (int): number of layers.
code: def __init__(
      self,
      hidden_size,
      num_attention_heads,
      attention_dropout,
      attention_mask_func,
      layer_number,
      apply_query_key_layer_scaling: bool = False,
      convert_fp16_to_fp32_in_softmax: bool = False,
      attn_mask_type=AttnMaskType.padding,
      masked_softmax_fusion=True,
      fp16=False,
      bf16=False,
  ):
      super().__init__()
      self.convert_fp16_to_fp32_in_softmax = (
          convert_fp16_to_fp32_in_softmax
      )
      self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
      self.attention_mask_func = attention_mask_func
      self.layer_number = layer_number
      self.hidden_size = hidden_size
      self.num_attention_heads = num_attention_heads
      self.attn_mask_type = attn_mask_type
      assert self.layer_number > 0
      self.attention_dropout = attention_dropout
  
      if self.apply_query_key_layer_scaling:
          self.convert_fp16_to_fp32_in_softmax = True
  
      assert (
          self.hidden_size % self.num_attention_heads == 0
      ), "hidden size is not divisible by the number of attention heads"
  
      self.hidden_size_per_attention_head = (
          self.hidden_size // num_attention_heads
      )
  
      self.world_size = gpc.get_world_size(ParallelMode.SEQUENCE)
  
      # Strided linear layer.
      self.query_key_value = _Linear(
          hidden_size,
          3 * self.hidden_size,
      )
  
      self.coeff = None
      self.norm_factor = math.sqrt(self.hidden_size)
  
      if self.apply_query_key_layer_scaling:
          self.coeff = layer_number
          self.norm_factor *= self.coeff
  
      self.scale_mask_softmax = FusedScaleMaskSoftmax(
          fp16,
          bf16,
          self.attn_mask_type,
          masked_softmax_fusion,
          self.attention_mask_func,
          self.convert_fp16_to_fp32_in_softmax,
          self.coeff,
      )
  
      self.attention_dropout = nn.Dropout(attention_dropout)
  
      # Output.
      self.dense = _Linear(
          hidden_size, hidden_size, bias=True, skip_bias_add=True
      )
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn import Parameter
  from colossalai.kernel import FusedScaleMaskSoftmax
  from colossalai.kernel.cuda_native.scaled_softmax import AttnMaskType
  from colossalai.legacy.context import seed
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.parallel_sequence._operation import RingAV, RingQK
  from colossalai.legacy.registry import LAYERS
label_desc: Use image node_id 1 to create a new TransformerSelfAttentionRing object from inherited base classes: nn with example: obj = TransformerSelfAttentionRing(hidden_size, num_attention_heads, attention_dropout, attention_mask_func, layer_number, apply_query_key_layer_scaling, convert_fp16_to_fp32_in_softmax, attn_mask_type, masked_softmax_fusion, fp16, bf16)