code: def __init__(
      self,
      embed_dim,
      num_heads,
      sparsity_config,
      bias=True,
      batch_first=True,
      attention_dropout=0.0,
      causal=False,
      max_seq_length=2048,
      device=None,
      dtype=None,
      **kwargs,
  ) -> None:
      assert batch_first
      factory_kwargs = {"device": device, "dtype": dtype}
      super().__init__()
      self.embed_dim = embed_dim
      self.causal = causal
  
      self.num_heads = num_heads
      assert (
          self.embed_dim % num_heads == 0
      ), "self.kdim must be divisible by num_heads"
      self.head_dim = self.embed_dim // num_heads
      assert self.head_dim in [
          16,
          32,
          64,
      ], "Only support head_dim == 16, 32, or 64"
  
      self.Wqkv = nn.Linear(
          embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
      )
      self.inner_attn = FlashBlocksparseAttention(
          sparsity_config,
          attention_dropout=attention_dropout,
          max_seq_length=max_seq_length,
          **factory_kwargs,
      )
      self.out_proj = nn.Linear(
          embed_dim, embed_dim, bias=bias, **factory_kwargs
      )
  
file_path: flash-attention/flash_attn/flash_blocksparse_attention.py
imports: import math
  import hydra
  import torch
  import torch.nn
  from einops import rearrange
  from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
  from flash_attn.flash_blocksparse_attn_interface import convert_blockmask, flash_blocksparse_attn_func
label_desc: Use image node_id 1 to create a new FlashBlocksparseMHA object from inherited base classes: nn with example: obj = FlashBlocksparseMHA(embed_dim, num_heads, sparsity_config, bias, batch_first, attention_dropout, causal, max_seq_length, device, dtype)