(self, x)
| 287 | return torch.einsum("bix,iox->box", x, weights) |
| 288 | |
| 289 | def forward(self, x): |
| 290 | B, N, c, k = x.shape # (B, N, c, k) |
| 291 | |
| 292 | x = x.view(B, N, -1) |
| 293 | x = x.permute(0, 2, 1) |
| 294 | x_fft = torch.fft.rfft(x) |
| 295 | # Multiply relevant Fourier modes |
| 296 | l = min(self.modes1, N // 2 + 1) |
| 297 | # l = N//2+1 |
| 298 | out_ft = torch.zeros(B, c * k, N // 2 + 1, device=x.device, dtype=torch.cfloat) |
| 299 | out_ft[:, :, :l] = self.compl_mul1d(x_fft[:, :, :l], self.weights1[:, :, :l]) |
| 300 | x = torch.fft.irfft(out_ft, n=N) |
| 301 | x = x.permute(0, 2, 1).view(B, N, c, k) |
| 302 | return x |
| 303 | |
| 304 | |
| 305 | # ## |
nothing calls this directly
no test coverage detected