class_docstr: A module which performs QKV attention. Matches legacy QKVAttention + input/ouput 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.
docstr_tok: ["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, encoder_kv=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
      )
      if encoder_kv is not None:
          assert encoder_kv.shape[1] == self.n_heads * ch * 2
          ek, ev = encoder_kv.reshape(
              bs * self.n_heads, ch * 2, -1
          ).split(ch, dim=1)
          k = th.cat([ek, k], dim=-1)
          v = th.cat([ev, v], dim=-1)
      scale = 1 / math.sqrt(math.sqrt(ch))
      weight = th.einsum(
          "bct,bcs->bts", q * scale, k * scale
      )  # 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)
      return a.reshape(bs, -1, length)
  
imports: import math
  from abc import abstractmethod
  import torch
  import torch.nn
  import torch.nn.functional
  from .fp16_util import convert_module_to_f16, convert_module_to_f32
  from .nn import avg_pool_nd, conv_nd, linear, normalization, timestep_embedding, zero_module
label_desc: Use image node_id 2 for calling the QKVAttention obj's underlying member method code with example usage: obj.forward(qkv, encoder_kv) and returns: a