(self, queries, keys, values, attn_mask)
| 53 | self.MWT_CZ = nn.ModuleList(MWT_CZ1d(k, alpha, L, c, base) for i in range(nCZ)) |
| 54 | |
| 55 | def forward(self, queries, keys, values, attn_mask): |
| 56 | B, L, H, E = queries.shape |
| 57 | _, S, _, D = values.shape |
| 58 | if L > S: |
| 59 | zeros = torch.zeros_like(queries[:, :(L - S), :]).float() |
| 60 | values = torch.cat([values, zeros], dim=1) |
| 61 | keys = torch.cat([keys, zeros], dim=1) |
| 62 | else: |
| 63 | values = values[:, :L, :, :] |
| 64 | keys = keys[:, :L, :, :] |
| 65 | values = values.view(B, L, -1) |
| 66 | |
| 67 | V = self.Lk0(values).view(B, L, self.c, -1) |
| 68 | for i in range(self.nCZ): |
| 69 | V = self.MWT_CZ[i](V) |
| 70 | if i < self.nCZ - 1: |
| 71 | V = F.relu(V) |
| 72 | |
| 73 | V = self.Lk1(V.view(B, L, -1)) |
| 74 | V = V.view(B, L, -1, D) |
| 75 | return (V.contiguous(), None) |
| 76 | |
| 77 | |
| 78 | class MultiWaveletCross(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected