| 457 | |
| 458 | class sparseKernelFT1d(nn.Module): |
| 459 | def __init__(self, |
| 460 | k, alpha, c=1, |
| 461 | nl=1, |
| 462 | initializer=None, |
| 463 | **kwargs): |
| 464 | super(sparseKernelFT1d, self).__init__() |
| 465 | |
| 466 | self.modes1 = alpha |
| 467 | self.scale = (1 / (c * k * c * k)) |
| 468 | self.weights1 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.float)) |
| 469 | self.weights2 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.float)) |
| 470 | self.weights1.requires_grad = True |
| 471 | self.weights2.requires_grad = True |
| 472 | self.k = k |
| 473 | |
| 474 | def compl_mul1d(self, order, x, weights): |
| 475 | x_flag = True |