docstr: Compute the windowed multi-head self-attention.
  
  Note: seq_length of keys and values (`seq_len_kv`) must be an integer
      multiple of the seq_length of the query (`seq_len_q`).
  
  Args:
      query: tensor of shape (b', seq_len_q, hidden_dim),
          representing the query
      key: tensor of shape (b', seq_len_kv, hidden_dim),
          representing the key
      value: tensor of shape (b', seq_len_kv, hidden_dim),
          representing the value
      mask: optional tensor of shape (b', seq_len_q, seq_len_q),
          representing the mask to apply to the attention; 1s will be masked
      -> b' is an augmented batch size that includes the total number of patches
      -> by default, hidden_dim is 768
  
  Returns:
      tensor of feautures of shape (b', seq_len_q, hidden_dim), as well as the
      attention matrix used, shape (b', num_heads, seq_len_q, seq_len_kv).
docstr_tok: ["Compute","the","windowed","multi-head","self-attention",".","Note",":","seq_length","of","keys","and","values","(","`","seq_len_kv","`",")","must","be","an","integer","multiple","of","the","seq_length","of","the","query","(","`","seq_len_q","`",")",".","Args",":","query",":","tensor","of","shape","(","b","'",",","seq_len_q",",","hidden_dim",")",",","representing","the","query","key",":","tensor","of","shape","(","b","'",",","seq_len_kv",",","hidden_dim",")",",","representing","the","key","value",":","tensor","of","shape","(","b","'",",","seq_len_kv",",","hidden_dim",")",",","representing","the","value","mask",":","optional","tensor","of","shape","(","b","'",",","seq_len_q",",","seq_len_q",")",",","representing","the","mask","to","apply","to","the","attention",";","1s","will","be","masked","-",">","b","'","is","an","augmented","batch","size","that","includes","the","total","number","of","patches","-",">","by","default",",","hidden_dim","is","768","Returns",":","tensor","of","feautures","of","shape","(","b","'",",","seq_len_q",",","hidden_dim",")",",","as","well","as","the","attention","matrix","used",",","shape","(","b","'",",","num_heads",",","seq_len_q",",","seq_len_kv",")","."]
code: def forward(
      self,
      query: Tensor,
      key: Tensor,
      value: Tensor,
      mask: Optional[torch.Tensor] = None,
  ) -> Tuple[Tensor, Tensor]:
      """Compute the windowed multi-head self-attention.
  
      Note: seq_length of keys and values (`seq_len_kv`) must be an integer
          multiple of the seq_length of the query (`seq_len_q`).
  
      Args:
          query: tensor of shape (b', seq_len_q, hidden_dim),
              representing the query
          key: tensor of shape (b', seq_len_kv, hidden_dim),
              representing the key
          value: tensor of shape (b', seq_len_kv, hidden_dim),
              representing the value
          mask: optional tensor of shape (b', seq_len_q, seq_len_q),
              representing the mask to apply to the attention; 1s will be masked
          -> b' is an augmented batch size that includes the total number of patches
          -> by default, hidden_dim is 768
  
      Returns:
          tensor of feautures of shape (b', seq_len_q, hidden_dim), as well as the
          attention matrix used, shape (b', num_heads, seq_len_q, seq_len_kv).
      """
      *b, seq_len_q, c = query.shape
      assert (
          c == self.hidden_dim
      ), f"Shape mismatch, {c} != {self.hidden_dim}"
      seq_len_kv = value.shape[-2]
      assert seq_len_kv % seq_len_q == 0, (
          f"seq_length of keys and values = {seq_len_kv}  must be an integer multiple "
          f"of the seq_length of the query = {seq_len_q}"
      )
      blowup = seq_len_kv // seq_len_q
  
      query = self.q_linear(query)  # [b', seq_len_q, hidden_dim]
      key = self.k_linear(key)  # [b', seq_len_kv, hidden_dim]
      value = self.v_linear(value)  # [b', seq_len_kv, hidden_dim]
  
      # reshape by splitting channels into num_heads, then permute:
      query = (
          query.reshape(
              *b, seq_len_q, self.num_heads, c // self.num_heads
          )
          .permute(0, 2, 1, 3)
          .contiguous()
      )  # [b', num_heads, seq_len_q, c // num_heads]
      key = (
          key.reshape(
              *b, seq_len_kv, self.num_heads, c // self.num_heads
          )
          .permute(0, 2, 1, 3)
          .contiguous()
      )  # [b', num_heads, seq_len_kv, c // num_heads]
      value = (
          value.reshape(
              *b, seq_len_kv, self.num_heads, c // self.num_heads
          )
          .permute(0, 2, 1, 3)
          .contiguous()
      )  # [b', num_heads, seq_len_kv, c // num_heads]
  
      # b', num_heads, seq_len_q, seq_len_kv
      attn = (
          torch.matmul(
              query,  # [b', num_heads, seq_len_q, c // num_heads]
              key.transpose(
                  2, 3
              ),  # [b', num_heads, c // num_heads, seq_len_kv]
          )  # [b', num_heads,seq_len_q, seq_len_kv]
          * self.attn_scale
      )
  
      if mask is not None:
          if mask.shape[-2:] != (seq_len_q, seq_len_q):
              # Note that we only mask for self-attention in the decoder,
              # where the attention matrix has shape (..., seq_len_q, seq_len_q).
              raise ValueError(f"Invalid mask shape: {mask.shape}.")
  
          # Here, we add the mask to the attention with a large negative multiplier,
          # as this goes into a softmax it will become 0
          tile_pattern = [1] * mask.dim()
          tile_pattern[-1] = blowup
          attn = attn + torch.tile(mask, tile_pattern) * -1e6
      else:
          tile_pattern = None
  
      attn = self.softmax(
          attn
      )  # [b', num_heads, seq_length_q, seq_length_kv]
  
      if mask is not None and tile_pattern is not None:
          # We use the mask again, to be ensure that no masked dimension
          # affects the output.
          keep = 1 - mask
          attn = attn * torch.tile(keep, tile_pattern)
  
      attn = self.attn_dropout(
          attn
      )  # [b', num_heads, seq_length_q, seq_length_kv]
  
      features = torch.matmul(
          attn, value
      )  # [b', num_heads, seq_len_q, d_model//num_heads], c=d_model
      assert features.shape == (
          *b,
          self.num_heads,
          seq_len_q,
          c // self.num_heads,
      )
  
      features = (
          features.permute(
              0, 2, 1, 3
          )  # switch num_heads, seq_len_q dimensions.
          .contiguous()
          .reshape(
              *b, -1, self.hidden_dim
          )  # flatten num_heads, seq_len_q dimensions
      )
      features = self.proj(features)
      features = self.proj_dropout(features)
      assert features.shape == (*b, seq_len_q, c)
      return features, attn
  
imports: from typing import Optional, Tuple
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 2 for calling the WindowMultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: features,  attn