class_docstr: This class implements the relative multihead implementation similar to that in Transformer XL
  https://arxiv.org/pdf/1901.02860.pdf
  
  Arguments
  ---------
  embed_dim : int
      Size of the encoder feature vectors from which keys and values are computed.
  num_heads: int
      Number of attention heads.
  dropout : float, optional
      Dropout rate.
  vbias: bool, optional
      Whether to use bias for computing value.
  vdim: int, optional
      Size for value. Default is embed_dim (Note each head is embed_dim // num_heads).
  mask_pos_future: bool, optional
      Whether to mask future positional encodings values.
      Must be true for causal applications e.g. decoder.
  Example
  -------
  >>> inputs = torch.rand([6, 60, 512])
  >>> pos_emb = torch.rand([1, 2*60-1, 512])
  >>> net = RelPosMHAXL(num_heads=8, embed_dim=inputs.shape[-1])
  >>> outputs, attn = net(inputs, inputs, inputs, pos_emb)
  >>> outputs.shape
  torch.Size([6, 60, 512])
docstr: Arguments
  ----------
  query : tensor
      (B, L, E) where L is the target sequence length,
      B is the batch size, E is the embedding dimension.
  key : tensor
      (B, S, E) where S is the source sequence length,
      B is the batch size, E is the embedding dimension.
  value : tensor
      (B, S, E) where S is the source sequence length,
      B is the batch size, E is the embedding dimension.
  pos_emb : tensor
      bidirectional sinusoidal positional embedding tensor (1, 2*S-1, E) where S is the max length between source and target sequence lengths,
      and E is the embedding dimension.
  key_padding_mask : tensor
      (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.
  attn_mask : tensor
      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.
  
  Outputs
  -------
  out : tensor
      (B, L, E) where L is the target sequence length, B is the
      batch size, E is the embedding dimension.
  attn_score : tensor
      (B, L, S) where B is the batch size, L is the target
      sequence length, S is the source sequence length.
docstr_tok: ["Arguments","--","--","--","--","--","query",":","tensor","(","B",",","L",",","E",")","where","L","is","the","target","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","key",":","tensor","(","B",",","S",",","E",")","where","S","is","the","source","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","value",":","tensor","(","B",",","S",",","E",")","where","S","is","the","source","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","pos_emb",":","tensor","bidirectional","sinusoidal","positional","embedding","tensor","(","1",",","2","*","S-1",",","E",")","where","S","is","the","max","length","between","source","and","target","sequence","lengths",",","and","E","is","the","embedding","dimension",".","key_padding_mask",":","tensor","(","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",".","attn_mask",":","tensor","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",".","Outputs","--","--","--","-","out",":","tensor","(","B",",","L",",","E",")","where","L","is","the","target","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","attn_score",":","tensor","(","B",",","L",",","S",")","where","B","is","the","batch","size",",","L","is","the","target","sequence","length",",","S","is","the","source","sequence","length","."]
code: def forward(
      self,
      query,
      key,
      value,
      pos_embs,
      key_padding_mask=None,
      attn_mask=None,
      return_attn_weights=True,
  ):
      """
      Arguments
      ----------
      query : tensor
          (B, L, E) where L is the target sequence length,
          B is the batch size, E is the embedding dimension.
      key : tensor
          (B, S, E) where S is the source sequence length,
          B is the batch size, E is the embedding dimension.
      value : tensor
          (B, S, E) where S is the source sequence length,
          B is the batch size, E is the embedding dimension.
      pos_emb : tensor
          bidirectional sinusoidal positional embedding tensor (1, 2*S-1, E) where S is the max length between source and target sequence lengths,
          and E is the embedding dimension.
      key_padding_mask : tensor
          (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.
      attn_mask : tensor
          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.
  
      Outputs
      -------
      out : tensor
          (B, L, E) where L is the target sequence length, B is the
          batch size, E is the embedding dimension.
      attn_score : tensor
          (B, L, S) where B is the batch size, L is the target
          sequence length, S is the source sequence length.
      """
  
      # query, key and value are of shape batch, time, embed_dim
      bsz = query.shape[0]
      klen = key.shape[1]
      qlen = query.shape[1]
  
      if self._qkv_same_embed_dim:
          # self-attention
          if (query is key or torch.equal(query, key)) and (
              key is value or torch.equal(key, value)
          ):
              query, key, value = (
                  nn.functional.linear(query, self.in_proj_weight)
                  .view(bsz, -1, self.num_heads, self.head_dim * 3)
                  .chunk(3, dim=-1)
              )
          else:
              qweight, kweight, vweight = self.in_proj_weight.chunk(
                  3, dim=0
              )
              query = nn.functional.linear(query, qweight).view(
                  bsz, -1, self.num_heads, self.head_dim
              )
              key = nn.functional.linear(key, kweight).view(
                  bsz, -1, self.num_heads, self.head_dim
              )
              value = nn.functional.linear(value, vweight).view(
                  bsz, -1, self.num_heads, self.head_dim
              )
      else:
          raise NotImplementedError
          query, key = (
              nn.functional.linear(query, self.qk_proj_weight)
              .view(bsz, -1, self.num_heads, self.head_dim * 2)
              .chunk(2, dim=-1)
          )
          value = nn.functional.linear(value, self.v_proj_weight).view(
              bsz, -1, self.num_heads, self.vhead_dim
          )
  
      if self.vbias is not None:
          value = value + self.value_bias_weight.view(
              1, 1, self.num_heads, self.vhead_dim
          )
  
      p_k = self.linear_pos(pos_embs).view(
          1, -1, self.num_heads, self.head_dim
      )
      # (batch, head, klen, d_k)
  
      q_with_bias_u = (
          query
          + self.pos_bias_u.view(1, 1, self.num_heads, self.head_dim)
      ).transpose(1, 2)
      # (batch, head, qlen, d_k)
      q_with_bias_v = (
          query
          + self.pos_bias_v.view(1, 1, self.num_heads, self.head_dim)
      ).transpose(1, 2)
  
      # Moved the `* self.scale` mul from after the `attn_score` sum to prior
      # to the matmul in order to lower overflow risks on fp16.
      # This change is inspired by the following paper, but no other changes
      # were ported from there so far.
      # ref: E.T.: Re-Thinking Self-Attention for Transformer Models on GPUs
      # https://asherliu.github.io/docs/sc21a.pdf
  
      # (batch, head, qlen, klen)
      matrix_ac = torch.matmul(
          q_with_bias_u * self.scale, key.permute(0, 2, 3, 1)
      )
      # (batch, num_heads, klen, 2*klen-1)
      matrix_bd = torch.matmul(
          q_with_bias_v * self.scale, p_k.permute(0, 2, 3, 1)
      )
      matrix_bd = self.rel_shift(matrix_bd)  # shifting trick
  
      # if klen != qlen:
      #   import ipdb
      #  ipdb.set_trace(
  
      attn_score = matrix_ac + matrix_bd  # already scaled above
  
      # compute attention probability
      if attn_mask is not None:
          if attn_mask.ndim == 2:
              attn_mask = attn_mask.view(1, 1, qlen, klen)
          else:
              attn_mask = attn_mask.view(-1, self.num_heads, qlen, klen)
  
          if attn_mask.dtype == torch.bool:
              attn_score = attn_score.masked_fill(
                  attn_mask, self.attn_fill_value
              )
          else:
              attn_score += attn_mask
  
      if key_padding_mask is not None:
          attn_score = attn_score.masked_fill(
              key_padding_mask.view(bsz, 1, 1, klen),
              self.attn_fill_value,
          )
  
      attn_score = F.softmax(attn_score, dim=-1, dtype=torch.float32)
      attn_score = self.dropout_att(attn_score)
  
      # it is possible for us to hit full NaN when using chunked training
      # so reapply masks, except with 0.0 instead as we are after the softmax
      # because -inf would output 0.0 regardless anyway
      if attn_mask is not None:
          if attn_mask.dtype == torch.bool:
              attn_score = attn_score.masked_fill(attn_mask, 0.0)
          else:
              # NOTE: the above fix is not implemented for this case as
              # summing the mask with NaN would still result in NaN
              pass
  
      if key_padding_mask is not None:
          attn_score = attn_score.masked_fill(
              key_padding_mask.view(bsz, 1, 1, klen),
              0.0,
          )
  
      x = torch.matmul(
          attn_score, value.transpose(1, 2)
      )  # (batch, head, time1, d_k)
      x = (
          x.transpose(1, 2)
          .contiguous()
          .view(bsz, -1, self.vhead_dim * self.num_heads)
      )  # (batch, time1, d_model)
  
      out = self.out_proj(x)
      if return_attn_weights:
          return out, attn_score
      return out
  
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 4 for calling the RelPosMHAXL obj's underlying member method code with example usage: obj.forward(query, key, value, pos_embs, key_padding_mask, attn_mask, return_attn_weights) and returns: out, out,  attn_score