class_docstr: Multi-headed self-attention for WavLM model :cite:`chen2022wavlm`.
  Wraps around ``torch.nn.MultiheadAttention``, creating relaive position embeddings and passing them to multi-headed
  attention as a mask.
  Source: https://github.com/microsoft/unilm/blob/2d8302f09c99bca2b82e6e868d81d4281cceebc8/wavlm/modules.py#L303-L763
  
  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: to ``0.0``)
      bias (bool, optional): If ``True``, add bias to input / output projection layers. (Default: ``True``)
      has_relative_attention_bias (bool, optional): If ``True``, apply relative position embedding.
          Necessary in the first encoder layer, but not in the subsequent ones. (Default: ``False``)
      num_buckets (int, optional): Number of buckets for relative position embedding. (Default: ``32``)
      max_distance (int, optional): Naximum distance for relative position embedding. (Default: ``128``)
      gru_rel_pos (bool, optional): If ``True``, apply gated relative position embedding. (Default: ``False``)
docstr: Args:
      query (Tensor): Input of shape ``(batch_size, src_len, embed_dim)``.
      key_padding_mask (Tensor or None, optional): Mask to exclude keys that are pads, of shape
          `(batch, src_len)`, where padding elements are indicated by 1s. (Default: ``None``)
      attn_mask: Needs to be ``None``. The argument exists for compatibility with
          ``EncoderLayer``. (Default: ``None``)
      position_bias (Tensor or None, optional): Position bias of shape
          ``(batch_size * num_heads, src_len, src_len)``. When used inside WavLM model encoder, will be
          generated in the first layer and then passed from each encoder layer to the next one.
          (Default: ``None``)
  Returns:
      attn_output (Tensor): Attention output of shape ``(batch_size, src_len, embed_dim)``.
      position_bias (Tensor or None): Position bias of shape ``(batch_size * num_heads, src_len, src_len)``.
code: def forward(
      self,
      query: Tensor,
      key_padding_mask: Optional[Tensor] = None,
      attention_mask: Optional[Tensor] = None,
      position_bias: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      """
      Args:
          query (Tensor): Input of shape ``(batch_size, src_len, embed_dim)``.
          key_padding_mask (Tensor or None, optional): Mask to exclude keys that are pads, of shape
              `(batch, src_len)`, where padding elements are indicated by 1s. (Default: ``None``)
          attn_mask: Needs to be ``None``. The argument exists for compatibility with
              ``EncoderLayer``. (Default: ``None``)
          position_bias (Tensor or None, optional): Position bias of shape
              ``(batch_size * num_heads, src_len, src_len)``. When used inside WavLM model encoder, will be
              generated in the first layer and then passed from each encoder layer to the next one.
              (Default: ``None``)
      Returns:
          attn_output (Tensor): Attention output of shape ``(batch_size, src_len, embed_dim)``.
          position_bias (Tensor or None): Position bias of shape ``(batch_size * num_heads, src_len, src_len)``.
      """
      bsz, seq_len, embed_dim = query.size()
      assert embed_dim == self.embed_dim
      assert attention_mask is None
  
      if self.rel_attn_embed is not None and position_bias is None:
          position_bias = self.compute_bias(seq_len, seq_len)
          position_bias = position_bias.unsqueeze(0).repeat(
              bsz, 1, 1, 1
          )
  
      attn_mask_rel_pos: Optional[Tensor] = None
      if position_bias is not None:
          attn_mask_rel_pos = position_bias
          if self.gru_rel_pos:  # Apply gating on relative position bias
              query_layer = query.view(bsz, seq_len, self.num_heads, -1)
              query_layer = query_layer.permute(0, 2, 1, 3)
  
              gate_a, gate_b = torch.sigmoid(
                  self.gru_rel_pos_linear(query_layer)
                  .view(bsz, self.num_heads, seq_len, 2, 4)
                  .sum(-1, keepdim=False)
              ).chunk(2, dim=-1)
              gate_a_1 = (
                  gate_a * (gate_b * self.gru_rel_pos_const - 1.0) + 2.0
              )
              attn_mask_rel_pos = (
                  gate_a_1.view(bsz, self.num_heads, -1, 1)
                  * position_bias
              )
  
          attn_mask_rel_pos = attn_mask_rel_pos.view(
              (bsz, self.num_heads, seq_len, seq_len)
          )
  
      if attn_mask_rel_pos is not None and key_padding_mask is not None:
          key_padding_mask = key_padding_mask.view(
              bsz, 1, 1, seq_len
          ).expand(-1, self.num_heads, -1, -1)
          key_padding_mask = torch.nn.functional._canonical_mask(
              mask=key_padding_mask,
              mask_name="key_padding_mask",
              other_type=torch.nn.functional._none_or_dtype(
                  attn_mask_rel_pos
              ),
              other_name="",
              target_type=query.dtype,
          )
      if attn_mask_rel_pos is not None and key_padding_mask is not None:
          attn_mask_rel_pos = attn_mask_rel_pos + key_padding_mask
      query_projected = torch.nn.functional.linear(
          query,
          self.attention.in_proj_weight,
          self.attention.in_proj_bias,
      )
      query, key, value = query_projected.chunk(3, -1)
      shape = (bsz, seq_len, self.num_heads, self.head_dim)
      query = query.view(shape).transpose(
          2, 1
      )  # (batch, num_heads, seq_len, head_dim)
      key = key.view(shape).transpose(
          2, 1
      )  # (batch, num_heads, seq_len, head_dim)
      value = value.view(shape).transpose(
          2, 1
      )  # (batch, num_heads, seq_len, head_dim)
      dropout = self.dropout if self.training else 0.0
      attn_output = torch.nn.functional.scaled_dot_product_attention(
          query,
          key,
          value,
          attn_mask=attn_mask_rel_pos,
          dropout_p=dropout,
          is_causal=False,
      )
      attn_output = attn_output.transpose(1, 2).reshape(
          bsz, -1, self.num_heads * self.head_dim
      )
      attn_output = self.attention.out_proj(attn_output)
      return attn_output, position_bias
  
file_path: audio/src/torchaudio/models/wav2vec2/wavlm_attention.py
imports: import math
  from typing import Optional, Tuple
  import torch
  from torch import nn, Tensor
label_desc: Use image node_id 4 for calling the WavLMSelfAttention obj's underlying member method code with example usage: obj.forward(query, key_padding_mask, attention_mask, position_bias) and returns: attn_output,  position_bias