(self, q, k, v, mask)
| 69 | return torch.einsum("bhi,hio->bho", input, weights) |
| 70 | |
| 71 | def forward(self, q, k, v, mask): |
| 72 | # size = [B, L, H, E] |
| 73 | B, L, H, E = q.shape |
| 74 | x = q.permute(0, 2, 3, 1) |
| 75 | # Compute Fourier coefficients |
| 76 | x_ft = torch.fft.rfft(x, dim=-1) |
| 77 | # Perform Fourier neural operations |
| 78 | out_ft = torch.zeros(B, H, E, L // 2 + 1, device=x.device, dtype=torch.cfloat) |
| 79 | for wi, i in enumerate(self.index): |
| 80 | out_ft[:, :, :, i] = self.compl_mul1d(x_ft[:, :, :, i], self.weights1[:, :, :, wi]) |
| 81 | # Return to time domain |
| 82 | x = torch.fft.irfft(out_ft, n=x.size(-1)) |
| 83 | return (x, None) |
| 84 | |
| 85 | |
| 86 | # ########## Fourier Cross Former #################### |
nothing calls this directly
no test coverage detected