class_docstr: A module which performs QKV attention. Matches legacy QKVAttention + input/output heads shaping
class_docstr_tok: ["A","module","which","performs","QKV","attention",".","Matches","legacy","QKVAttention","+","input\/output","heads","shaping"]
docstr: Apply QKV attention.
  
  :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
  :return: an [N x (H * C) x T] tensor after attention.
code: def forward(self, qkv, mask=None, rel_pos=None):
      """
      Apply QKV attention.
  
      :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
      :return: 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.reshape(bs * self.n_heads, ch * 3, length).split(
          ch, dim=1
      )
      scale = 1 / math.sqrt(math.sqrt(ch))
      weight = torch.einsum(
          "bct,bcs->bts", q * scale, k * scale
      )  # More stable with f16 than dividing afterwards
      if rel_pos is not None:
          weight = rel_pos(
              weight.reshape(
                  bs, self.n_heads, weight.shape[-2], weight.shape[-1]
              )
          ).reshape(
              bs * self.n_heads, weight.shape[-2], weight.shape[-1]
          )
      weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
      if mask is not None:
          # The proper way to do this is to mask before the softmax using -inf, but that doesn't work properly on CPUs.
          mask = mask.repeat(self.n_heads, 1).unsqueeze(1)
          weight = weight * mask
      a = torch.einsum("bts,bcs->bct", weight, v)
  
      return a.reshape(bs, -1, length)
  
imports: import functools
  import math
  import os
  import fsspec
  import torch
  import torch.nn
  import torch.nn.functional
  import torchaudio
  from transformers import LogitsWarper
  from TTS.tts.layers.tortoise.xtransformers import ContinuousTransformerWrapper, RelativePositionBias
label_desc: Use image node_id 2 for calling the QKVAttentionLegacy obj's underlying member method code with example usage: obj.forward(qkv, mask, rel_pos) and returns: a