class_docstr: Cross-attention layer class
  
  Cross-attention layer takes input with size [s, b, h] and context with size
  [s, b, h] and returns output of the same size.
class_docstr_tok: ["Cross-attention","layer","class","Cross-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","context","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
code: def __init__(
      self,
      config: TransformerConfig,
      submodules: CrossAttentionSubmodules,
      layer_number: int,
      attn_mask_type=AttnMaskType.padding,
  ):
      super().__init__(
          config=config,
          submodules=submodules,
          layer_number=layer_number,
          attn_mask_type=attn_mask_type,
          attention_type="cross",
      )
  
      if (
          self.config.num_query_groups
          != self.config.num_attention_heads
      ):
          raise ValueError(
              f"Group query attention is not currently supported in cross attention."
          )
      assert self.query_projection_size == self.kv_projection_size
  
      self.linear_q = build_module(
          submodules.linear_q,
          self.config.hidden_size,
          self.query_projection_size,
          config=self.config,
          init_method=self.config.init_method,
          gather_output=False,
          bias=self.config.add_bias_linear,
          skip_bias_add=False,
          is_expert=False,
      )
  
      self.linear_kv = build_module(
          submodules.linear_kv,
          self.config.hidden_size,
          2 * self.kv_projection_size,
          config=self.config,
          init_method=self.config.init_method,
          gather_output=False,
          bias=self.config.add_bias_linear,
          skip_bias_add=False,
          is_expert=False,
      )
  
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 CrossAttention object from inherited base classes: Attention with example: obj = CrossAttention(config, submodules, layer_number, attn_mask_type)