code: def __init__(
      self,
      dim: int,
      head_dim: int,
      n_heads: int,
      n_kv_heads: int,
      rope_theta: float,
  ):
      super().__init__()
      mp_size = mp_utils.get_world_size()
  
      self.head_dim = head_dim
      self.rope_theta = rope_theta
  
      self.n_local_heads = n_heads // mp_size
      self.n_local_kv_heads = n_kv_heads // mp_size
  
      self.wqkv = nn.Linear(
          dim,
          (self.n_local_heads + 2 * self.n_local_kv_heads) * head_dim,
          bias=False,
      )
      self.wo = nn.Linear(
          self.n_local_heads * head_dim,
          dim,
          bias=False,
      )
      self._register_load_state_dict_pre_hook(self.load_hook)
  
imports: from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import mp_utils
  import torch
  from torch import nn
  from torch.nn import functional
  from xformers.ops import RMSNorm, fmha, rope_padded
  from xformers.ops.fmha.attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(dim, head_dim, n_heads, n_kv_heads, rope_theta)