class_docstr: A module which performs QKV attention and splits in a different order.
class_docstr_tok: ["A","module","which","performs","QKV","attention","and","splits","in","a","different","order","."]
docstr: Apply QKV attention.
  :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
  :return: an [N x (H * C) x T] tensor after attention.
docstr_tok: ["Apply","QKV","attention",".",":","param","qkv",":","an","[","N","x","(","3","*","H","*","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: th.Tensor) -> th.Tensor:
      """
      Apply QKV attention.
      :param qkv: an [N x (3 * H * 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.chunk(3, dim=1)
      scale = 1 / math.sqrt(math.sqrt(ch))
      weight = th.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 = th.softmax(weight.float(), dim=-1).type(weight.dtype)
      a = th.einsum(
          "bts,bcs->bct",
          weight,
          v.reshape(bs * self.n_heads, ch, length),
      )
      return a.reshape(bs, -1, length)
  
imports: import logging
  import math
  from abc import abstractmethod
  from typing import Iterable, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from torch.utils.checkpoint import checkpoint
  from ...modules.attention import SpatialTransformer
  from ...modules.diffusionmodules.util import avg_pool_nd, conv_nd, linear, normalization, timestep_embedding, zero_module
  from ...modules.video_attention import SpatialVideoTransformer
  from ...util import exists
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