| 270 | |
| 271 | class sparseKernelFT1d(nn.Module): |
| 272 | def __init__(self, |
| 273 | k, alpha, c=1, |
| 274 | nl=1, |
| 275 | initializer=None, |
| 276 | **kwargs): |
| 277 | super(sparseKernelFT1d, self).__init__() |
| 278 | |
| 279 | self.modes1 = alpha |
| 280 | self.scale = (1 / (c * k * c * k)) |
| 281 | self.weights1 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.cfloat)) |
| 282 | self.weights1.requires_grad = True |
| 283 | self.k = k |
| 284 | |
| 285 | def compl_mul1d(self, x, weights): |
| 286 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) |