class_docstr: Parallel chunked cross-attention layer class.
  
  Self-attention layer takes input with size [b, s, h]
  and returns output of the same size.
class_docstr_tok: ["Parallel","chunked","cross-attention","layer","class",".","Self-attention","layer","takes","input","with","size","[","b",",","s",",","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,
      precision=16,
      apply_query_key_layer_scaling=False,
      kv_channels=None,
      masked_softmax_fusion=True,
      attention_dropout=0.1,
      megatron_legacy=False,
      chunk_size=64,  # each chunk, how many tokens
      bias=True,
      headscale=False,
      normalize_attention_scores=True,
  ):
      super(ParallelChunkedCrossAttention, self).__init__(config=config)
      self.cross_attention = ParallelAttention(
          config=config,
          init_method=init_method,
          output_layer_init_method=output_layer_init_method,
          layer_number=layer_number,
          num_attention_heads=num_attention_heads,
          hidden_size=hidden_size,
          attention_type=AttnType.cross_attn,
          attn_mask_type=AttnMaskType.padding,
          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,
          megatron_legacy=megatron_legacy,
          bias=bias,
          headscale=headscale,
          normalize_attention_scores=normalize_attention_scores,
      )
      self.chunk_size = chunk_size
  
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 ParallelChunkedCrossAttention object from inherited base classes: MegatronModule with example: obj = ParallelChunkedCrossAttention(config, init_method, output_layer_init_method, layer_number, num_attention_heads, hidden_size, precision, apply_query_key_layer_scaling, kv_channels, masked_softmax_fusion, attention_dropout, megatron_legacy, chunk_size, bias, headscale, normalize_attention_scores)