code: def forward(
      self,
      Q: torch.Tensor,
      K: Optional[torch.Tensor] = None,
      V: Optional[torch.Tensor] = None,
      prev: Optional[torch.Tensor] = None,
      key_padding_mask: Optional[torch.Tensor] = None,
      attn_mask: Optional[torch.Tensor] = None,
  ):
      bs = Q.size(0)
      if K is None:
          K = Q
      if V is None:
          V = Q
  
      # Linear (+ split in multiple heads)
      q_s = (
          self.W_Q(Q)
          .view(bs, -1, self.n_heads, self.d_k)
          .transpose(1, 2)
      )  # q_s    : [bs x n_heads x max_q_len x d_k]
      k_s = (
          self.W_K(K)
          .view(bs, -1, self.n_heads, self.d_k)
          .permute(0, 2, 3, 1)
      )  # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
      v_s = (
          self.W_V(V)
          .view(bs, -1, self.n_heads, self.d_v)
          .transpose(1, 2)
      )  # v_s    : [bs x n_heads x q_len x d_v]
  
      # Apply Scaled Dot-Product Attention (multiple heads)
      if self.res_attention:
          output, attn_weights, attn_scores = self.sdp_attn(
              q_s,
              k_s,
              v_s,
              prev=prev,
              key_padding_mask=key_padding_mask,
              attn_mask=attn_mask,
          )
      else:
          output, attn_weights = self.sdp_attn(
              q_s,
              k_s,
              v_s,
              key_padding_mask=key_padding_mask,
              attn_mask=attn_mask,
          )
      # output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]
  
      # back to the original inputs dimensions
      output = (
          output.transpose(1, 2)
          .contiguous()
          .view(bs, -1, self.n_heads * self.d_v)
      )  # output: [bs x q_len x n_heads * d_v]
      output = self.to_out(output)
  
      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 _MultiheadAttention 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