docstr: Windowed multi-head self-attention
  
  Args:
      hidden_dim: size of the hidden units
      num_heads: number of attention heads
      attn_drop: dropout rate of the attention layer. Defaults to 0.0.
      proj_drop: dropout rate of the projection layer. Defaults to 0.0.
  
  Raises:
      ValueError: if `hidden_dim` is not a multiple of `num_heads`
docstr_tok: ["Windowed","multi-head","self-attention","Args",":","hidden_dim",":","size","of","the","hidden","units","num_heads",":","number","of","attention","heads","attn_drop",":","dropout","rate","of","the","attention","layer",".","Defaults","to","0.0.","proj_drop",":","dropout","rate","of","the","projection","layer",".","Defaults","to","0.0",".","Raises",":","ValueError",":","if","`","hidden_dim","`","is","not","a","multiple","of","`","num_heads","`"]
code: def __init__(
      self,
      hidden_dim: int,
      num_heads: int,
      attn_drop: float = 0.0,
      proj_drop: float = 0.0,
  ) -> None:
      """Windowed multi-head self-attention
  
      Args:
          hidden_dim: size of the hidden units
          num_heads: number of attention heads
          attn_drop: dropout rate of the attention layer. Defaults to 0.0.
          proj_drop: dropout rate of the projection layer. Defaults to 0.0.
  
      Raises:
          ValueError: if `hidden_dim` is not a multiple of `num_heads`
      """
      super().__init__()
      self.hidden_dim = hidden_dim
      self.num_heads = num_heads
      if hidden_dim % num_heads != 0:
          raise ValueError(
              f"Size of hidden units ({hidden_dim}) not divisible by number "
              f"of heads ({num_heads})."
          )
      # constnat to scale output before softmax-ing
      self.attn_scale = (hidden_dim // num_heads) ** (-0.5)
      # All linear layers are initialized with truncated normal
      self.q_linear = torch.nn.Linear(
          in_features=hidden_dim, out_features=hidden_dim, bias=True
      )
      init_weights_truncated_normal(self.q_linear)
      self.k_linear = torch.nn.Linear(
          in_features=hidden_dim, out_features=hidden_dim, bias=True
      )
      init_weights_truncated_normal(self.k_linear)
      self.v_linear = torch.nn.Linear(
          in_features=hidden_dim, out_features=hidden_dim, bias=True
      )
      init_weights_truncated_normal(self.v_linear)
      self.attn_dropout = nn.Dropout(p=attn_drop)
      self.proj = torch.nn.Linear(
          in_features=hidden_dim, out_features=hidden_dim
      )
      init_weights_truncated_normal(self.proj)
      self.proj_dropout = torch.nn.Dropout(p=proj_drop)
      self.softmax = torch.nn.Softmax(dim=-1)
  
imports: from typing import Optional, Tuple
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 1 to create a new WindowMultiHeadAttention object from inherited base classes: nn with example: obj = WindowMultiHeadAttention(hidden_dim, num_heads, attn_drop, proj_drop)