x: [B, L, C].
(self, x)
| 74 | self.proj = nn.Linear(dim, dim) |
| 75 | |
| 76 | def forward(self, x): |
| 77 | """ |
| 78 | x: [B, L, C]. |
| 79 | """ |
| 80 | b, s, c, n, d = *x.size(), self.num_heads, self.head_dim |
| 81 | |
| 82 | # compute query, key, value |
| 83 | q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2) |
| 84 | |
| 85 | # compute attention |
| 86 | p = self.attn_dropout if self.training else 0.0 |
| 87 | x = attention(q, k, v, dropout_p=p, causal=self.causal, attention_type="none") |
| 88 | x = x.reshape(b, s, c) |
| 89 | |
| 90 | # output |
| 91 | x = self.proj(x) |
| 92 | x = F.dropout(x, self.proj_dropout, self.training) |
| 93 | return x |
| 94 | |
| 95 | |
| 96 | class SwiGLU(nn.Module): |
nothing calls this directly
no test coverage detected