class_docstr: The class is a wrapper of MultiHead Attention for torch.nn.MultiHeadAttention.
  
  Reference: https://pytorch.org/docs/stable/nn.html
  
  Arguments
  ----------
  num_heads : int
      parallel attention heads.
  dropout : float
      a Dropout layer on attn_output_weights (default: 0.0).
  bias : bool
      add bias as module parameter (default: True).
  add_bias_kv : bool
      add bias to the key and value sequences at dim=0.
  add_zero_attn : bool
      add a new batch of zeros to the key and value sequences at dim=1.
  kdim : int
      total number of features in key (default: None).
  vdim : int
      total number of features in value (default: None).
  
  Example
  -------
  >>> inputs = torch.rand([8, 60, 512])
  >>> net = MultiheadAttention(nhead=8, d_model=inputs.shape[-1])
  >>> outputs, attn = net(inputs, inputs, inputs)
  >>> outputs.shape
  torch.Size([8, 60, 512])
docstr: Arguments
  ----------
  query : torch.Tensor
      (B, L, E) where L is the target sequence length,
      B is the batch size, E is the embedding dimension.
  key : torch.Tensor
      (B, S, E) where S is the source sequence length,
      B is the batch size, E is the embedding dimension.
  value : torch.Tensor
      (B, S, E) where S is the source sequence length,
      B is the batch size, E is the embedding dimension.
  attn_mask : torch.Tensor, optional
      2D mask (L, S) where L is the target sequence length, S is
      the source sequence length.
      3D mask (N*num_heads, L, S) where N is the batch
      size, L is the target sequence length, S is the source sequence
      length. attn_mask ensure that position i is allowed to attend the
      unmasked positions. If a ByteTensor is provided, the non-zero
      positions are not allowed to attend while the zero positions will
      be unchanged. If a BoolTensor is provided, positions with True is
      not allowed to attend while False values will be unchanged. If a
      FloatTensor is provided, it will be added to the attention weight.
  key_padding_mask : torch.Tensor, optional
      (B, S) where B is the batch size, S is the source sequence
      length. If a ByteTensor is provided, the non-zero positions will
      be ignored while the position with the zero positions will be
      unchanged. If a BoolTensor is provided, the positions with the
      value of True will be ignored while the position with the value
      of False will be unchanged.
  return_attn_weights : bool, optional
      True to additionally return the attention weights, False otherwise.
  pos_embs : torch.Tensor, optional
      Positional embeddings added to the attention map of shape (L, S, E) or (L, S, 1).
  
  Outputs
  -------
  attn_output : torch.Tensor
      (B, L, E) where L is the target sequence length, B is the
      batch size, E is the embedding dimension.
  attn_output_weights : torch.Tensor
      (B, L, S) where B is the batch size, L is the target
      sequence length, S is the source sequence length.
      This is returned only if `return_attn_weights=True` (True by default).
docstr_tok: ["Arguments","--","--","--","--","--","query",":","torch.Tensor","(","B",",","L",",","E",")","where","L","is","the","target","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","key",":","torch.Tensor","(","B",",","S",",","E",")","where","S","is","the","source","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","value",":","torch.Tensor","(","B",",","S",",","E",")","where","S","is","the","source","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","attn_mask",":","torch.Tensor",",","optional","2D","mask","(","L",",","S",")","where","L","is","the","target","sequence","length",",","S","is","the","source","sequence","length",".","3D","mask","(","N","*","num_heads",",","L",",","S",")","where","N","is","the","batch","size",",","L","is","the","target","sequence","length",",","S","is","the","source","sequence","length",".","attn_mask","ensure","that","position","i","is","allowed","to","attend","the","unmasked","positions",".","If","a","ByteTensor","is","provided",",","the","non-zero","positions","are","not","allowed","to","attend","while","the","zero","positions","will","be","unchanged",".","If","a","BoolTensor","is","provided",",","positions","with","True","is","not","allowed","to","attend","while","False","values","will","be","unchanged",".","If","a","FloatTensor","is","provided",",","it","will","be","added","to","the","attention","weight",".","key_padding_mask",":","torch.Tensor",",","optional","(","B",",","S",")","where","B","is","the","batch","size",",","S","is","the","source","sequence","length",".","If","a","ByteTensor","is","provided",",","the","non-zero","positions","will","be","ignored","while","the","position","with","the","zero","positions","will","be","unchanged",".","If","a","BoolTensor","is","provided",",","the","positions","with","the","value","of","True","will","be","ignored","while","the","position","with","the","value","of","False","will","be","unchanged",".","return_attn_weights",":","bool",",","optional","True","to","additionally","return","the","attention","weights",",","False","otherwise",".","pos_embs",":","torch.Tensor",",","optional","Positional","embeddings","added","to","the","attention","map","of","shape","(","L",",","S",",","E",")","or","(","L",",","S",",","1",")",".","Outputs","--","--","--","-","attn_output",":","torch.Tensor","(","B",",","L",",","E",")","where","L","is","the","target","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","attn_output_weights",":","torch.Tensor","(","B",",","L",",","S",")","where","B","is","the","batch","size",",","L","is","the","target","sequence","length",",","S","is","the","source","sequence","length",".","This","is","returned","only","if","`","return_attn_weights=True","`","(","True","by","default",")","."]
code: def forward(
      self,
      query,
      key,
      value,
      attn_mask: Optional[torch.Tensor] = None,
      key_padding_mask: Optional[torch.Tensor] = None,
      return_attn_weights: bool = True,
      pos_embs: Optional[torch.Tensor] = None,
  ):
      """
      Arguments
      ----------
      query : torch.Tensor
          (B, L, E) where L is the target sequence length,
          B is the batch size, E is the embedding dimension.
      key : torch.Tensor
          (B, S, E) where S is the source sequence length,
          B is the batch size, E is the embedding dimension.
      value : torch.Tensor
          (B, S, E) where S is the source sequence length,
          B is the batch size, E is the embedding dimension.
      attn_mask : torch.Tensor, optional
          2D mask (L, S) where L is the target sequence length, S is
          the source sequence length.
          3D mask (N*num_heads, L, S) where N is the batch
          size, L is the target sequence length, S is the source sequence
          length. attn_mask ensure that position i is allowed to attend the
          unmasked positions. If a ByteTensor is provided, the non-zero
          positions are not allowed to attend while the zero positions will
          be unchanged. If a BoolTensor is provided, positions with True is
          not allowed to attend while False values will be unchanged. If a
          FloatTensor is provided, it will be added to the attention weight.
      key_padding_mask : torch.Tensor, optional
          (B, S) where B is the batch size, S is the source sequence
          length. If a ByteTensor is provided, the non-zero positions will
          be ignored while the position with the zero positions will be
          unchanged. If a BoolTensor is provided, the positions with the
          value of True will be ignored while the position with the value
          of False will be unchanged.
      return_attn_weights : bool, optional
          True to additionally return the attention weights, False otherwise.
      pos_embs : torch.Tensor, optional
          Positional embeddings added to the attention map of shape (L, S, E) or (L, S, 1).
  
      Outputs
      -------
      attn_output : torch.Tensor
          (B, L, E) where L is the target sequence length, B is the
          batch size, E is the embedding dimension.
      attn_output_weights : torch.Tensor
          (B, L, S) where B is the batch size, L is the target
          sequence length, S is the source sequence length.
          This is returned only if `return_attn_weights=True` (True by default).
      """
      # give tensors of shape (time, batch, fea)
      query = query.permute(1, 0, 2)
      key = key.permute(1, 0, 2)
      value = value.permute(1, 0, 2)
  
      # this will be legit because of https://github.com/pytorch/pytorch/blob/5288d05cfdda85c46c4df84617fa7f37c21b10b3/torch/nn/functional.py#L4946
      # we can inject relative learnable pos embeddings directly in MHA via the attn_mask
      if pos_embs is not None:
          if attn_mask is not None:
              attn_mask += pos_embs
          else:
              attn_mask = pos_embs
  
      output, attention_weights = self.att(
          query,
          key,
          value,
          attn_mask=attn_mask,
          key_padding_mask=key_padding_mask,
          need_weights=return_attn_weights,
      )
  
      # reshape the output back to (batch, time, fea)
      output = output.permute(1, 0, 2)
  
      if return_attn_weights:
          return output, attention_weights
  
      return output
  
imports: import torch
  import logging
  import torch.nn
  import numpy
  from typing import Optional
  from speechbrain.dataio.dataio import length_to_mask
  import torch.nn.functional
  import math
label_desc: Use image node_id 2 for calling the MultiheadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, attn_mask, key_padding_mask, return_attn_weights, pos_embs) and returns: output, output,  attention_weights