| 317 | |
| 318 | # FourierGNN |
| 319 | def fourierGC(self, x, B, N, L): |
| 320 | o1_real = torch.zeros([B, (N*L)//2 + 1, self.frequency_size * self.hidden_size_factor], |
| 321 | device=x.device) |
| 322 | o1_imag = torch.zeros([B, (N*L)//2 + 1, self.frequency_size * self.hidden_size_factor], |
| 323 | device=x.device) |
| 324 | o2_real = torch.zeros(x.shape, device=x.device) |
| 325 | o2_imag = torch.zeros(x.shape, device=x.device) |
| 326 | |
| 327 | o3_real = torch.zeros(x.shape, device=x.device) |
| 328 | o3_imag = torch.zeros(x.shape, device=x.device) |
| 329 | |
| 330 | o1_real = F.relu( |
| 331 | torch.einsum('bli,ii->bli', x.real, self.w1[0]) - \ |
| 332 | torch.einsum('bli,ii->bli', x.imag, self.w1[1]) + \ |
| 333 | self.b1[0] |
| 334 | ) |
| 335 | |
| 336 | o1_imag = F.relu( |
| 337 | torch.einsum('bli,ii->bli', x.imag, self.w1[0]) + \ |
| 338 | torch.einsum('bli,ii->bli', x.real, self.w1[1]) + \ |
| 339 | self.b1[1] |
| 340 | ) |
| 341 | |
| 342 | # 1 layer |
| 343 | y = torch.stack([o1_real, o1_imag], dim=-1) |
| 344 | y = F.softshrink(y, lambd=self.sparsity_threshold) |
| 345 | |
| 346 | o2_real = F.relu( |
| 347 | torch.einsum('bli,ii->bli', o1_real, self.w2[0]) - \ |
| 348 | torch.einsum('bli,ii->bli', o1_imag, self.w2[1]) + \ |
| 349 | self.b2[0] |
| 350 | ) |
| 351 | |
| 352 | o2_imag = F.relu( |
| 353 | torch.einsum('bli,ii->bli', o1_imag, self.w2[0]) + \ |
| 354 | torch.einsum('bli,ii->bli', o1_real, self.w2[1]) + \ |
| 355 | self.b2[1] |
| 356 | ) |
| 357 | |
| 358 | # 2 layer |
| 359 | x = torch.stack([o2_real, o2_imag], dim=-1) |
| 360 | x = F.softshrink(x, lambd=self.sparsity_threshold) |
| 361 | x = x + y |
| 362 | |
| 363 | o3_real = F.relu( |
| 364 | torch.einsum('bli,ii->bli', o2_real, self.w3[0]) - \ |
| 365 | torch.einsum('bli,ii->bli', o2_imag, self.w3[1]) + \ |
| 366 | self.b3[0] |
| 367 | ) |
| 368 | |
| 369 | o3_imag = F.relu( |
| 370 | torch.einsum('bli,ii->bli', o2_imag, self.w3[0]) + \ |
| 371 | torch.einsum('bli,ii->bli', o2_real, self.w3[1]) + \ |
| 372 | self.b3[1] |
| 373 | ) |
| 374 | |
| 375 | # 3 layer |
| 376 | z = torch.stack([o3_real, o3_imag], dim=-1) |