| 69 | return d_u, d_v |
| 70 | |
| 71 | class TriangularToeplitzMultFast(torch.autograd.Function): |
| 72 | @staticmethod |
| 73 | def forward(ctx, u, v): |
| 74 | n = u.shape[-1] |
| 75 | u_expand = F.pad(u, (0, n)) |
| 76 | v_expand = F.pad(v, (0, n)) |
| 77 | u_f = torch.fft.rfft(u_expand, n=2*n, dim=-1) |
| 78 | v_f = torch.fft.rfft(v_expand, n=2*n, dim=-1) |
| 79 | |
| 80 | ctx.save_for_backward(u_f, v_f) |
| 81 | |
| 82 | uv_f = u_f * v_f |
| 83 | output = torch.fft.irfft(uv_f, n=2*n, dim=-1)[..., :n] |
| 84 | return output |
| 85 | |
| 86 | @staticmethod |
| 87 | def backward(ctx, grad): |
| 88 | u_f, v_f = ctx.saved_tensors |
| 89 | n = grad.shape[-1] |
| 90 | g_expand = F.pad(grad.flip(-1), (0, n)) |
| 91 | g_f = torch.fft.rfft(g_expand, n=2*n, dim=-1) |
| 92 | gu_f = g_f * u_f |
| 93 | gv_f = g_f * v_f |
| 94 | d_u = torch.fft.irfft(gv_f, n=2*n, dim=-1)[..., :n] |
| 95 | d_v = torch.fft.irfft(gu_f, n=2*n, dim=-1)[..., :n] |
| 96 | d_u = d_u.flip(-1) |
| 97 | d_v = d_v.flip(-1) |
| 98 | return d_u, d_v |
| 99 | |
| 100 | class TriangularToeplitzMultPadded(torch.autograd.Function): |
| 101 | @staticmethod |
nothing calls this directly
no outgoing calls
no test coverage detected