class_docstr: A module which performs QKV attention and splits in a different order.
  
  Example
  -------
  >>> attn = QKVAttention(4)
  >>> n = 4
  >>> c = 8
  >>> h = 64
  >>> w = 16
  >>> qkv = torch.randn(4, (3 * h * c), w)
  >>> out = attn(qkv)
  >>> out.shape
  torch.Size([4, 512, 16])
class_docstr_tok: ["A","module","which","performs","QKV","attention","and","splits","in","a","different","order",".","Example","--","--","--","-",">",">",">","attn","=","QKVAttention","(","4",")",">",">",">","n","=","4",">",">",">","c","=","8",">",">",">","h","=","64",">",">",">","w","=","16",">",">",">","qkv","=","torch.randn","(","4",",","(","3","*","h","*","c",")",",","w",")",">",">",">","out","=","attn","(","qkv",")",">",">",">","out.shape","torch.Size","(","[","4",",","512",",","16","]",")"]
docstr: Apply QKV attention.
  
  Arguments
  ---------
  qkv: torch.Tensor
      an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
  
  Results
  -------
  result: torch.Tensor
      an [N x (H * C) x T] tensor after attention.
docstr_tok: ["Apply","QKV","attention",".","Arguments","--","--","--","--","-","qkv",":","torch.Tensor","an","[","N","x","(","3","*","H","*","C",")","x","T","]","tensor","of","Qs",",","Ks",",","and","Vs",".","Results","--","--","--","-","result",":","torch.Tensor","an","[","N","x","(","H","*","C",")","x","T","]","tensor","after","attention","."]
code: def forward(self, qkv):
      """Apply QKV attention.
  
      Arguments
      ---------
      qkv: torch.Tensor
          an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
  
      Results
      -------
      result: torch.Tensor
          an [N x (H * C) x T] tensor after attention.
      """
      bs, width, length = qkv.shape
      assert width % (3 * self.n_heads) == 0
      ch = width // (3 * self.n_heads)
      q, k, v = qkv.chunk(3, dim=1)
      scale = 1 / math.sqrt(math.sqrt(ch))
      weight = torch.einsum(
          "bct,bcs->bts",
          (q * scale).view(bs * self.n_heads, ch, length),
          (k * scale).view(bs * self.n_heads, ch, length),
      )  # More stable with f16 than dividing afterwards
      weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
      a = torch.einsum(
          "bts,bcs->bct",
          weight,
          v.reshape(bs * self.n_heads, ch, length),
      )
      return a.reshape(bs, -1, length)
  
imports: from abc import abstractmethod
  from speechbrain.utils.data_utils import pad_divisible
  from .autoencoders import NormalizingAutoencoder
  import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the QKVAttention obj's underlying member method code with example usage: obj.forward(qkv) and returns: a