def forward(self, qkv, y):
    '\n        Apply QKV attention.\n\n        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.\n        :return: an [N x (H * C) x T] tensor after attention.\n        '
    (bs, width, length) = qkv.shape
    if (y is None):
        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)
    else:
        assert ((width % self.n_heads) == 0)
        ch = (width // self.n_heads)
        q = qkv.reshape((bs * self.n_heads), ch, length)
        k = v = y.reshape((bs * self.n_heads), ch, (- 1))
    scale = (1 / math.sqrt(math.sqrt(ch)))
    weight = th.einsum('bct,bcs->bts', (q * scale), (k * scale))
    weight = self.dropout(th.softmax(weight.float(), dim=(- 1)).type(weight.dtype))
    a = th.einsum('bts,bcs->bct', weight, v)
    return a.reshape(bs, (- 1), length)
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.
def forward(self, qkv, y):
    '\n        Apply QKV attention.\n\n        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.\n        :return: an [N x (H * C) x T] tensor after attention.\n        '
    (bs, width, length) = qkv.shape
    if (y is None):
        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)
    else:
        assert ((width % self.n_heads) == 0)
        ch = (width // self.n_heads)
        q = qkv.reshape((bs * self.n_heads), ch, length)
        k = v = y.reshape((bs * self.n_heads), ch, (- 1))
    scale = (1 / math.sqrt(math.sqrt(ch)))
    weight = th.einsum('bct,bcs->bts', (q * scale), (k * scale))
    weight = self.dropout(th.softmax(weight.float(), dim=(- 1)).type(weight.dtype))
    a = th.einsum('bts,bcs->bct', weight, v)
    return a.reshape(bs, (- 1), length)