class_docstr: ### Cross Attention Layer
  
  This falls-back to self-attention when conditional embeddings are not specified.
class_docstr_tok: ["#","#","#","Cross","Attention","Layer","This","falls-back","to","self-attention","when","conditional","embeddings","are","not","specified","."]
docstr: #### Flash Attention
  
  :param q: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
  :param k: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
  :param v: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
docstr_tok: ["#","#","#","#","Flash","Attention",":","param","q",":","are","the","query","vectors","before","splitting","heads",",","of","shape","`","[","batch_size",",","seq",",","d_attn","]","`",":","param","k",":","are","the","query","vectors","before","splitting","heads",",","of","shape","`","[","batch_size",",","seq",",","d_attn","]","`",":","param","v",":","are","the","query","vectors","before","splitting","heads",",","of","shape","`","[","batch_size",",","seq",",","d_attn","]","`"]
code: def flash_attention(
      self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
  ):
      """
      #### Flash Attention
  
      :param q: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
      :param k: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
      :param v: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
      """
  
      # Get batch size and number of elements along sequence axis (`width * height`)
      batch_size, seq_len, _ = q.shape
  
      # Stack `q`, `k`, `v` vectors for flash attention, to get a single tensor of
      # shape `[batch_size, seq_len, 3, n_heads * d_head]`
      qkv = torch.stack((q, k, v), dim=2)
      # Split the heads
      qkv = qkv.view(batch_size, seq_len, 3, self.n_heads, self.d_head)
  
      # Flash attention works for head sizes `32`, `64` and `128`, so we have to pad the heads to
      # fit this size.
      if self.d_head <= 32:
          pad = 32 - self.d_head
      elif self.d_head <= 64:
          pad = 64 - self.d_head
      elif self.d_head <= 128:
          pad = 128 - self.d_head
      else:
          raise ValueError(
              f"Head size ${self.d_head} too large for Flash Attention"
          )
  
      # Pad the heads
      if pad:
          qkv = torch.cat(
              (
                  qkv,
                  qkv.new_zeros(
                      batch_size, seq_len, 3, self.n_heads, pad
                  ),
              ),
              dim=-1,
          )
  
      # Compute attention
      # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)V$$
      # This gives a tensor of shape `[batch_size, seq_len, n_heads, d_padded]`
      out, _ = self.flash(qkv)
      # Truncate the extra head size
      out = out[:, :, :, : self.d_head]
      # Reshape to `[batch_size, seq_len, n_heads * d_head]`
      out = out.reshape(batch_size, seq_len, self.n_heads * self.d_head)
  
      # Map to `[batch_size, height * width, d_model]` with a linear layer
      return self.to_out(out)
  
imports: from typing import Optional
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 3 for calling the CrossAttention obj's underlying member method code with example usage: obj.flash_attention(q, k, v) and returns: self