-# Y[:, 0] = A[:, 0] * Y0 + X[:,0]
-# for t > 0 Y[:, t] = A[:, t] * Y[:, t - 1] + X[:, t]
-
-
-def pscan_rec(A, X, Y0):
- if X.size(1) % 2 == 1:
- if X.size(1) == 1:
- return A[:, :1] * Y0[:, None] + X[:, :1]
- else:
- Y = pscan_rec(A[:, :-1], X[:, :-1], Y0)
- return torch.cat([Y, A[:, -1:] * Y[:, -1:] + X[:, -1:]], dim=1)
-
- A2 = A.reshape(A.size(0), A.size(1) // 2, 2, A.size(2))
- X2 = X.reshape(X.size(0), X.size(1) // 2, 2, X.size(2))
-
- X_star = X2[:, :, 0].clone()
- X_star[:, 1:] += A2[:, 1:, 0] * X2[:, :-1, 1]
-
- A_star = A2[:, :, 0].clone()
- A_star[:, 1:] *= A2[:, :-1, 1]
-
- Y_star = pscan_rec(A_star, X_star, Y0)[:, :, None]
-
- Y = torch.cat([Y_star, A2[:, :, 1, None] * Y_star + X2[:, :, 1, None]], dim=2)
-
- Y = Y.reshape(Y.size(0), -1, Y.size(-1))
+# can be computed as
+#
+# Y[:, t] = A[:, t] * Y0 + X[:, t]
+
+
+def expand(A, X):
+ if A.size(1) == 1:
+ return
+ T = 2 * (A.size(1) // 2)
+ Aa = A[:, :T].view(A.size(0), T // 2, 2, -1)
+ Xa = X[:, :T].view(X.size(0), T // 2, 2, -1)
+ Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+ Aa[:, :, 1].mul_(Aa[:, :, 0])
+ expand(Aa[:, :, 1], Xa[:, :, 1])
+ Xa[:, 1:, 0].add_(Aa[:, 1:, 0].mul(Xa[:, :-1, 1]))
+ Aa[:, 1:, 0].mul_(Aa[:, :-1, 1])
+ if T < A.size(1):
+ X[:, -1].add_(A[:, -1].mul(X[:, -2]))
+ A[:, -1].mul_(A[:, -2])
+
+
+# Computes inplace Y[:, s] = \sum_{t >= s} X[:, t]
+
+
+def accrev(X):
+ if X.size(1) == 1:
+ return
+ T = 2 * (X.size(1) // 2)
+ Xa = X[:, -T:].view(X.size(0), T // 2, 2, -1)
+ Xa[:, :, 0].add_(Xa[:, :, 1])
+ accrev(Xa[:, :, 0])
+ Xa[:, :-1, 1].add_(Xa[:, 1:, 0])
+ if T < X.size(1):
+ X[:, 0].add_(X[:, 1])
+
+
+class PScan(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, A, X, Y0):
+ ctx.A = A[:, :, None].clone()
+ ctx.Y0 = Y0[:, None, :].clone()
+ ctx.A_star = A[:, :, None].clone()
+ ctx.X_star = X.clone()
+ expand(ctx.A_star, ctx.X_star)
+ return ctx.A_star * ctx.Y0 + ctx.X_star
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ U = grad_output * ctx.A_star
+ R = U.clone()
+ accrev(R)
+ Q = ctx.Y0 / ctx.A
+ Q[:, 1:].add_(ctx.X_star[:, :-1] / ctx.A_star[:, 1:])
+ return (Q * R).sum(-1), R / ctx.A_star, U
+
+
+pscan = PScan.apply