x: [B, L, C].
(self, x, mask)
| 29 | self.dropout = nn.Dropout(dropout) |
| 30 | |
| 31 | def forward(self, x, mask): |
| 32 | """ |
| 33 | x: [B, L, C]. |
| 34 | """ |
| 35 | b, s, c, n, d = *x.size(), self.num_heads, self.head_dim |
| 36 | |
| 37 | # compute query, key, value |
| 38 | q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3) |
| 39 | k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3) |
| 40 | v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3) |
| 41 | |
| 42 | # compute attention |
| 43 | p = self.dropout.p if self.training else 0.0 |
| 44 | x = F.scaled_dot_product_attention(q, k, v, mask, p) |
| 45 | x = x.permute(0, 2, 1, 3).reshape(b, s, c) |
| 46 | |
| 47 | # output |
| 48 | x = self.o(x) |
| 49 | x = self.dropout(x) |
| 50 | return x |
| 51 | |
| 52 | |
| 53 | class AttentionBlock(nn.Module): |
nothing calls this directly
no test coverage detected