class_docstr: Multihead Self Attention module
  
  Args:
      embed_dim (int): Total dimension of the model.
      num_heads (int): The number of heads.
      dropout (float, optional):
          Dropout probability on attn_output_weights. Default: ``0.0``
class_docstr_tok: ["Multihead","Self","Attention","module","Args",":","embed_dim","(","int",")",":","Total","dimension","of","the","model",".","num_heads","(","int",")",":","The","number","of","heads",".","dropout","(","float",",","optional",")",":","Dropout","probability","on","attn_output_weights",".","Default",":","``","0.0","``"]
docstr: Args:
      x (Tensor): shape: ``[batch_size, sequence_length, embed_dim]``.
      attention_mask (Tensor or ``None``, optional):
          shape: ``[batch_size, 1, sequence_length, sequence_length]``
      position_bias: Not used. Only for the compatibility with :py:class:`WavLMSelfAttention`.
      key_padding_mask (Tensor or ``None``): Not used. Only for the compatibility with
          :py:class:`WavLMSelfAttention`.
  Returns:
      (Tensor, ``None``): The resulting attention output and ``None`` (necessary for compatibility
          with :py:class:`WavLMSelAttention`).
          Attention output shape: ``[batch, sequence_length, embed_dim]``.
docstr_tok: ["Args",":","x","(","Tensor",")",":","shape",":","``","[","batch_size",",","sequence_length",",","embed_dim","]","``",".","attention_mask","(","Tensor","or","``","None","``",",","optional",")",":","shape",":","``","[","batch_size",",","1",",","sequence_length",",","sequence_length","]","``","position_bias",":","Not","used",".","Only","for","the","compatibility","with",":","py",":","class",":","`","WavLMSelfAttention","`",".","key_padding_mask","(","Tensor","or","``","None","``",")",":","Not","used",".","Only","for","the","compatibility","with",":","py",":","class",":","`","WavLMSelfAttention","`",".","Returns",":","(","Tensor",",","``","None","``",")",":","The","resulting","attention","output","and","``","None","``","(","necessary","for","compatibility","with",":","py",":","class",":","`","WavLMSelAttention","`",")",".","Attention","output","shape",":","``","[","batch",",","sequence_length",",","embed_dim","]","``","."]
code: def forward(
      self,
      x: Tensor,
      attention_mask: Optional[Tensor] = None,
      position_bias: Optional[Tensor] = None,
      key_padding_mask: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      """
      Args:
          x (Tensor): shape: ``[batch_size, sequence_length, embed_dim]``.
          attention_mask (Tensor or ``None``, optional):
              shape: ``[batch_size, 1, sequence_length, sequence_length]``
          position_bias: Not used. Only for the compatibility with :py:class:`WavLMSelfAttention`.
          key_padding_mask (Tensor or ``None``): Not used. Only for the compatibility with
              :py:class:`WavLMSelfAttention`.
      Returns:
          (Tensor, ``None``): The resulting attention output and ``None`` (necessary for compatibility
              with :py:class:`WavLMSelAttention`).
              Attention output shape: ``[batch, sequence_length, embed_dim]``.
      """
      if x.ndim != 3 or x.shape[2] != self.embed_dim:
          raise ValueError(
              f"The expected input shape is (batch, sequence, embed_dim=={self.embed_dim}). "
              f"Found {x.shape}."
          )
      batch_size, length, embed_dim = x.size()
      if attention_mask is not None:
          shape_ = (batch_size, 1, length, length)
          if attention_mask.size() != shape_:
              raise ValueError(
                  f"The expected attention mask shape is {shape_}. "
                  f"Found {attention_mask.size()}."
              )
  
      shape = (batch_size, length, self.num_heads, self.head_dim)
      q = self.q_proj(x).view(*shape).transpose(2, 1)  # B, nH, L, Hd
      k = self.k_proj(x).view(*shape).transpose(2, 1)  # B, nH, L, Hd
      v = self.v_proj(x).view(*shape).transpose(2, 1)  # B, nH, L, Hd
      dropout = self.dropout if self.training else 0.0
      attn_output = torch.nn.functional.scaled_dot_product_attention(
          q,
          k,
          v,
          attn_mask=attention_mask,
          dropout_p=dropout,
          is_causal=False,
      )
      attn_output = attn_output.transpose(1, 2).reshape(
          batch_size, -1, self.num_heads * self.head_dim
      )
      output = self.out_proj(attn_output)
      return output, None
  
imports: import logging
  from typing import List, Optional, Tuple
  import torch
  from torch import nn, Tensor
  from torch.nn import Module, Parameter
  from .wavlm_attention import WavLMSelfAttention
label_desc: Use image node_id 2 for calling the SelfAttention obj's underlying member method code with example usage: obj.forward(x, attention_mask, position_bias, key_padding_mask) and returns: output,  None