x: B, T, D
(self, x, emb, src_mask)
| 206 | self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout) |
| 207 | |
| 208 | def forward(self, x, emb, src_mask): |
| 209 | """ |
| 210 | x: B, T, D |
| 211 | """ |
| 212 | B, T, D = x.shape |
| 213 | H = self.num_head |
| 214 | # B, T, 1, D |
| 215 | query = self.query(self.norm(x)).unsqueeze(2) |
| 216 | # B, 1, T, D |
| 217 | key = self.key(self.norm(x)).unsqueeze(1) |
| 218 | query = query.view(B, T, H, -1) |
| 219 | key = key.view(B, T, H, -1) |
| 220 | # B, T, T, H |
| 221 | attention = torch.einsum('bnhd,bmhd->bnmh', query, key) / math.sqrt(D // H) |
| 222 | attention = attention + (1 - src_mask.unsqueeze(-1)) * -100000 |
| 223 | weight = self.dropout(F.softmax(attention, dim=2)) |
| 224 | value = self.value(self.norm(x)).view(B, T, H, -1) |
| 225 | y = torch.einsum('bnmh,bmhd->bnhd', weight, value).reshape(B, T, D) |
| 226 | y = x + self.proj_out(y, emb) |
| 227 | return y |
| 228 | |
| 229 | class TemporalCrossAttention(nn.Module): |
| 230 |
nothing calls this directly
no outgoing calls
no test coverage detected