docstr: Multi Head Attention Layer
  Input shape:
      Q:       [batch_size (bs) x max_q_len x hidden_size]
      K, V:    [batch_size (bs) x q_len x hidden_size]
      mask:    [q_len x q_len]
docstr_tok: ["Multi","Head","Attention","Layer","Input","shape",":","Q",":","[","batch_size","(","bs",")","x","max_q_len","x","hidden_size","]","K",",","V",":","[","batch_size","(","bs",")","x","q_len","x","hidden_size","]","mask",":","[","q_len","x","q_len","]"]
code: def __init__(
      self,
      hidden_size,
      n_heads,
      d_k=None,
      d_v=None,
      res_attention=False,
      attn_dropout=0.0,
      proj_dropout=0.0,
      qkv_bias=True,
      lsa=False,
  ):
      """
      Multi Head Attention Layer
      Input shape:
          Q:       [batch_size (bs) x max_q_len x hidden_size]
          K, V:    [batch_size (bs) x q_len x hidden_size]
          mask:    [q_len x q_len]
      """
      super().__init__()
      d_k = hidden_size // n_heads if d_k is None else d_k
      d_v = hidden_size // n_heads if d_v is None else d_v
  
      self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
  
      self.W_Q = nn.Linear(hidden_size, d_k * n_heads, bias=qkv_bias)
      self.W_K = nn.Linear(hidden_size, d_k * n_heads, bias=qkv_bias)
      self.W_V = nn.Linear(hidden_size, d_v * n_heads, bias=qkv_bias)
  
      # Scaled Dot-Product Attention (multiple heads)
      self.res_attention = res_attention
      self.sdp_attn = _ScaledDotProductAttention(
          hidden_size,
          n_heads,
          attn_dropout=attn_dropout,
          res_attention=self.res_attention,
          lsa=lsa,
      )
  
      # Poject output
      self.to_out = nn.Sequential(
          nn.Linear(n_heads * d_v, hidden_size),
          nn.Dropout(proj_dropout),
      )
  
imports: import math
  import numpy
  from typing import Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from ..common._base_windows import BaseWindows
  from ..losses.pytorch import MAE
label_desc: Use image node_id 1 to create a new _MultiheadAttention object from inherited base classes: nn with example: obj = _MultiheadAttention(hidden_size, n_heads, d_k, d_v, res_attention, attn_dropout, proj_dropout, qkv_bias, lsa)