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.
(self, qkv)
| 357 | self.n_heads = n_heads |
| 358 | |
| 359 | def forward(self, qkv): |
| 360 | """ |
| 361 | Apply QKV attention. |
| 362 | :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. |
| 363 | :return: an [N x (H * C) x T] tensor after attention. |
| 364 | """ |
| 365 | bs, width, length = qkv.shape |
| 366 | assert width % (3 * self.n_heads) == 0 |
| 367 | ch = width // (3 * self.n_heads) |
| 368 | q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1) |
| 369 | scale = 1 / math.sqrt(math.sqrt(ch)) |
| 370 | weight = th.einsum( |
| 371 | "bct,bcs->bts", q * scale, k * scale |
| 372 | ) # More stable with f16 than dividing afterwards |
| 373 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) |
| 374 | a = th.einsum("bts,bcs->bct", weight, v) |
| 375 | return a.reshape(bs, -1, length) |
| 376 | |
| 377 | @staticmethod |
| 378 | def count_flops(model, _x, y): |
nothing calls this directly
no outgoing calls
no test coverage detected