| import math |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
|
|
| def npo2(len): |
| |
|
|
| return 2 ** math.ceil(math.log2(len)) |
|
|
| def pad_npo2(X): |
| |
|
|
| len_npo2 = npo2(X.size(1)) |
| pad_tuple = (0, 0, 0, 0, 0, len_npo2 - X.size(1)) |
| return F.pad(X, pad_tuple, "constant", 0) |
|
|
| class PScan(torch.autograd.Function): |
| @staticmethod |
| def pscan(A, X): |
| |
| |
| B, D, L, _ = A.size() |
| num_steps = int(math.log2(L)) |
|
|
| |
| Aa = A |
| Xa = X |
| for _ in range(num_steps-2): |
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T//2, 2, -1) |
| Xa = Xa.view(B, D, T//2, 2, -1) |
| |
| Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0])) |
| Aa[:, :, :, 1].mul_(Aa[:, :, :, 0]) |
|
|
| Aa = Aa[:, :, :, 1] |
| Xa = Xa[:, :, :, 1] |
|
|
| |
| if Xa.size(2) == 4: |
| Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0])) |
| Aa[:, :, 1].mul_(Aa[:, :, 0]) |
|
|
| Xa[:, :, 3].add_(Aa[:, :, 3].mul(Xa[:, :, 2] + Aa[:, :, 2].mul(Xa[:, :, 1]))) |
| elif Xa.size(2) == 2: |
| Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0])) |
| return |
| else: |
| return |
|
|
| |
| Aa = A[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)] |
| Xa = X[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)] |
| Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 1])) |
| Aa[:, :, 2].mul_(Aa[:, :, 1]) |
|
|
| for k in range(num_steps-3, -1, -1): |
| Aa = A[:, :, 2**k-1:L:2**k] |
| Xa = X[:, :, 2**k-1:L:2**k] |
|
|
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T//2, 2, -1) |
| Xa = Xa.view(B, D, T//2, 2, -1) |
|
|
| Xa[:, :, 1:, 0].add_(Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1])) |
| Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1]) |
|
|
| @staticmethod |
| def pscan_rev(A, X): |
| |
|
|
| B, D, L, _ = A.size() |
| num_steps = int(math.log2(L)) |
|
|
| |
| Aa = A |
| Xa = X |
| for _ in range(num_steps-2): |
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T//2, 2, -1) |
| Xa = Xa.view(B, D, T//2, 2, -1) |
| |
| Xa[:, :, :, 0].add_(Aa[:, :, :, 0].mul(Xa[:, :, :, 1])) |
| Aa[:, :, :, 0].mul_(Aa[:, :, :, 1]) |
|
|
| Aa = Aa[:, :, :, 0] |
| Xa = Xa[:, :, :, 0] |
|
|
| |
| if Xa.size(2) == 4: |
| Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 3])) |
| Aa[:, :, 2].mul_(Aa[:, :, 3]) |
|
|
| Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1].add(Aa[:, :, 1].mul(Xa[:, :, 2])))) |
| elif Xa.size(2) == 2: |
| Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1])) |
| return |
| else: |
| return |
|
|
| |
| Aa = A[:, :, 0:L:2**(num_steps-2)] |
| Xa = X[:, :, 0:L:2**(num_steps-2)] |
| Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 2])) |
| Aa[:, :, 1].mul_(Aa[:, :, 2]) |
|
|
| for k in range(num_steps-3, -1, -1): |
| Aa = A[:, :, 0:L:2**k] |
| Xa = X[:, :, 0:L:2**k] |
|
|
| T = Xa.size(2) |
| Aa = Aa.view(B, D, T//2, 2, -1) |
| Xa = Xa.view(B, D, T//2, 2, -1) |
|
|
| Xa[:, :, :-1, 1].add_(Aa[:, :, :-1, 1].mul(Xa[:, :, 1:, 0])) |
| Aa[:, :, :-1, 1].mul_(Aa[:, :, 1:, 0]) |
|
|
| @staticmethod |
| def forward(ctx, A_in, X_in): |
| |
|
|
| L = X_in.size(1) |
|
|
| |
| if L == npo2(L): |
| A = A_in.clone() |
| X = X_in.clone() |
| else: |
| |
| A = pad_npo2(A_in) |
| X = pad_npo2(X_in) |
| |
| |
| A = A.transpose(2, 1) |
| X = X.transpose(2, 1) |
|
|
| |
| PScan.pscan(A, X) |
|
|
| ctx.save_for_backward(A_in, X) |
| |
| |
| return X.transpose(2, 1)[:, :L] |
| |
| @staticmethod |
| def backward(ctx, grad_output_in): |
| |
|
|
| A_in, X = ctx.saved_tensors |
|
|
| L = grad_output_in.size(1) |
|
|
| |
| if L == npo2(L): |
| grad_output = grad_output_in.clone() |
| |
| else: |
| grad_output = pad_npo2(grad_output_in) |
| A_in = pad_npo2(A_in) |
|
|
| |
| grad_output = grad_output.transpose(2, 1) |
| A_in = A_in.transpose(2, 1) |
| A = torch.nn.functional.pad(A_in[:, :, 1:], (0, 0, 0, 1)) |
|
|
| |
| PScan.pscan_rev(A, grad_output) |
|
|
| Q = torch.zeros_like(X) |
| Q[:, :, 1:].add_(X[:, :, :-1] * grad_output[:, :, 1:]) |
|
|
| return Q.transpose(2, 1)[:, :L], grad_output.transpose(2, 1)[:, :L] |
| |
| pscan = PScan.apply |
|
|
