x: [B, L1, C]. context: [B, L2, C] or None. mask: [B, L2] or [B, L1, L2] or None.
(self, x, context=None, mask=None, pos_bias=None)
| 73 | self.dropout = nn.Dropout(dropout) |
| 74 | |
| 75 | def forward(self, x, context=None, mask=None, pos_bias=None): |
| 76 | """ |
| 77 | x: [B, L1, C]. |
| 78 | context: [B, L2, C] or None. |
| 79 | mask: [B, L2] or [B, L1, L2] or None. |
| 80 | """ |
| 81 | # check inputs |
| 82 | context = x if context is None else context |
| 83 | b, n, c = x.size(0), self.num_heads, self.head_dim |
| 84 | |
| 85 | # compute query, key, value |
| 86 | q = self.q(x).view(b, -1, n, c) |
| 87 | k = self.k(context).view(b, -1, n, c) |
| 88 | v = self.v(context).view(b, -1, n, c) |
| 89 | |
| 90 | # attention bias |
| 91 | attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) |
| 92 | if pos_bias is not None: |
| 93 | attn_bias += pos_bias |
| 94 | if mask is not None: |
| 95 | assert mask.ndim in [2, 3] |
| 96 | mask = mask.view(b, 1, 1, |
| 97 | -1) if mask.ndim == 2 else mask.unsqueeze(1) |
| 98 | attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) |
| 99 | |
| 100 | # compute attention (T5 does not use scaling) |
| 101 | attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias |
| 102 | attn = F.softmax(attn.float(), dim=-1).type_as(attn) |
| 103 | x = torch.einsum('bnij,bjnc->binc', attn, v) |
| 104 | |
| 105 | # output |
| 106 | x = x.reshape(b, -1, n * c) |
| 107 | x = self.o(x) |
| 108 | x = self.dropout(x) |
| 109 | return x |
| 110 | |
| 111 | |
| 112 | class T5FeedForward(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected