class_docstr: Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
  (Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
  by Lee et al, 2021)
class_docstr_tok: ["Scaled","Dot-Product","Attention","module","(","Attention","is","all","you","need","by","Vaswani","et","al.",",","2017",")","with","optional","residual","attention","from","previous","layer","(","Realformer",":","Transformer","likes","residual","attention","by","He","et","al",",","2020",")","and","locality","self","sttention","(","Vision","Transformer","for","Small-Size","Datasets","by","Lee","et","al",",","2021",")"]
docstr: Input shape:
      q               : [bs x n_heads x max_q_len x d_k]
      k               : [bs x n_heads x d_k x seq_len]
      v               : [bs x n_heads x seq_len x d_v]
      prev            : [bs x n_heads x q_len x seq_len]
      key_padding_mask: [bs x seq_len]
      attn_mask       : [1 x seq_len x seq_len]
  Output shape:
      output:  [bs x n_heads x q_len x d_v]
      attn   : [bs x n_heads x q_len x seq_len]
      scores : [bs x n_heads x q_len x seq_len]
docstr_tok: ["Input","shape",":","q",":","[","bs","x","n_heads","x","max_q_len","x","d_k","]","k",":","[","bs","x","n_heads","x","d_k","x","seq_len","]","v",":","[","bs","x","n_heads","x","seq_len","x","d_v","]","prev",":","[","bs","x","n_heads","x","q_len","x","seq_len","]","key_padding_mask",":","[","bs","x","seq_len","]","attn_mask",":","[","1","x","seq_len","x","seq_len","]","Output","shape",":","output",":","[","bs","x","n_heads","x","q_len","x","d_v","]","attn",":","[","bs","x","n_heads","x","q_len","x","seq_len","]","scores",":","[","bs","x","n_heads","x","q_len","x","seq_len","]"]
code: def forward(
      self,
      q: torch.Tensor,
      k: torch.Tensor,
      v: torch.Tensor,
      prev: Optional[torch.Tensor] = None,
      key_padding_mask: Optional[torch.Tensor] = None,
      attn_mask: Optional[torch.Tensor] = None,
  ):
      """
      Input shape:
          q               : [bs x n_heads x max_q_len x d_k]
          k               : [bs x n_heads x d_k x seq_len]
          v               : [bs x n_heads x seq_len x d_v]
          prev            : [bs x n_heads x q_len x seq_len]
          key_padding_mask: [bs x seq_len]
          attn_mask       : [1 x seq_len x seq_len]
      Output shape:
          output:  [bs x n_heads x q_len x d_v]
          attn   : [bs x n_heads x q_len x seq_len]
          scores : [bs x n_heads x q_len x seq_len]
      """
  
      # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
      attn_scores = (
          torch.matmul(q, k) * self.scale
      )  # attn_scores : [bs x n_heads x max_q_len x q_len]
  
      # Add pre-softmax attention scores from the previous layer (optional)
      if prev is not None:
          attn_scores = attn_scores + prev
  
      # Attention mask (optional)
      if (
          attn_mask is not None
      ):  # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
          if attn_mask.dtype == torch.bool:
              attn_scores.masked_fill_(attn_mask, -np.inf)
          else:
              attn_scores += attn_mask
  
      # Key padding mask (optional)
      if (
          key_padding_mask is not None
      ):  # mask with shape [bs x q_len] (only when max_w_len == q_len)
          attn_scores.masked_fill_(
              key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf
          )
  
      # normalize the attention weights
      attn_weights = F.softmax(
          attn_scores, dim=-1
      )  # attn_weights   : [bs x n_heads x max_q_len x q_len]
      attn_weights = self.attn_dropout(attn_weights)
  
      # compute the new values given the attention weights
      output = torch.matmul(
          attn_weights, v
      )  # output: [bs x n_heads x max_q_len x d_v]
  
      if self.res_attention:
          return output, attn_weights, attn_scores
      else:
          return output, attn_weights
  
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 2 for calling the _ScaledDotProductAttention obj's underlying member method code with example usage: obj.forward(q, k, v, prev, key_padding_mask, attn_mask) and returns: output,  attn_weights,  attn_scores, output,  attn_weights