code: def blocksparse_attention(
      q: torch.Tensor,
      k: torch.Tensor,
      v: torch.Tensor,
      dropout: Optional[torch.nn.Module] = None,
      block_size: int = 128,
  ) -> torch.Tensor:
      orig_dim = q.dim()
      seq_len = q.shape[-2]
      # Layout head dimension: 1 or batch size (q.shape[0])
      layout_heads = 1
  
      # TODO perhaps add functionality to pad qkv if sequence length is not divisible by block size?
      assert (
          seq_len % block_size == 0
      ), "Sequence length must be divisible by block size"
  
      if orig_dim == 3:
          # Reshape from (N, S, hs) to (B, nh, S, hs) where N = B x nh, hs = D / nh
          # Assuming num_heads = 1, (N, S, hs) to (B, 1, S, hs)
          if layout_heads == 1:
              q = q.unsqueeze(1)
              k = k.unsqueeze(1)
              v = v.unsqueeze(1)
          else:
              q = q.unsqueeze(0)
              k = k.unsqueeze(0)
              v = v.unsqueeze(0)
  
      blocksparse_attention = _retrieve_blocksparse(
          layout_heads, seq_len, block_size
      )
      # Dropout is a no-op in evaluation mode
      if isinstance(dropout, torch.nn.Dropout):
          blocksparse_attention.attn_drop = dropout
      else:
          blocksparse_attention.attn_drop = torch.nn.Dropout(0.0)
      att = blocksparse_attention(q, k, v)
  
      # Reshape attention (B, nh, S, hs) back to (N, S, hs)
      if orig_dim == 3:
          return att.flatten(0, 1)
      return att
  
file_path: xformers/xformers/components/attention/core.py
imports: import logging
  import math
  from contextlib import nullcontext
  from functools import lru_cache
  from typing import Optional, Union
  import torch
  from xformers import _has_cpp_library, _is_triton_available
  from xformers.components.attention.attention_mask import AttentionMask
label_desc: Use image node_id 8 for calling a global function with example usage: blocksparse_attention(q, k, v, dropout, block_size) and returns: att, att