pscan.py

   1 #!/usr/bin/env python
   2
   3 import torch
   4
   5 ######################################################################
   6
   7 # Given A is NxTx1 and X is NxTxD, expands A and X in place in O(T),
   8 # and O(log(T)) if not core-bounded, so that
   9 #
  10 # Y[:, 0] = Y0
  11 # Y[:, t] = A[:, t] * Y[:, t-1] + X[:, t]
  12 #
  13 # can be computed as
  14 #
  15 # Y[:, t] = A[:, t] * Y0 + X[:, t]
  16
  17
  18 def expand(A, X):
  19     if A.size(1) == 1:
  20         return
  21     T = 2 * (A.size(1) // 2)
  22     Aa = A[:, :T].view(A.size(0), T // 2, 2, -1)
  23     Xa = X[:, :T].view(X.size(0), T // 2, 2, -1)
  24     Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
  25     Aa[:, :, 1].mul_(Aa[:, :, 0])
  26     expand(Aa[:, :, 1], Xa[:, :, 1])
  27     Xa[:, 1:, 0].add_(Aa[:, 1:, 0].mul(Xa[:, :-1, 1]))
  28     Aa[:, 1:, 0].mul_(Aa[:, :-1, 1])
  29     if T < A.size(1):
  30         X[:, -1].add_(A[:, -1].mul(X[:, -2]))
  31         A[:, -1].mul_(A[:, -2])
  32
  33
  34 # Computes inplace Y[:, s] = \sum_{t >= s} X[:, t]
  35
  36
  37 def accrev(X):
  38     if X.size(1) == 1:
  39         return
  40     T = 2 * (X.size(1) // 2)
  41     Xa = X[:, -T:].view(X.size(0), T // 2, 2, -1)
  42     Xa[:, :, 0].add_(Xa[:, :, 1])
  43     accrev(Xa[:, :, 0])
  44     Xa[:, :-1, 1].add_(Xa[:, 1:, 0])
  45     if T < X.size(1):
  46         X[:, 0].add_(X[:, 1])
  47
  48
  49 class PScan(torch.autograd.Function):
  50     @staticmethod
  51     def forward(ctx, A, X, Y0):
  52         ctx.A = A[:, :, None].clone()
  53         ctx.Y0 = Y0[:, None, :].clone()
  54         ctx.A_star = A[:, :, None].clone()
  55         ctx.X_star = X.clone()
  56         expand(ctx.A_star, ctx.X_star)
  57         return ctx.A_star * ctx.Y0 + ctx.X_star
  58
  59     @staticmethod
  60     def backward(ctx, grad_output):
  61         U = grad_output * ctx.A_star
  62         R = U.clone()
  63         accrev(R)
  64         Q = ctx.Y0 / ctx.A
  65         Q[:, 1:].add_(ctx.X_star[:, :-1] / ctx.A_star[:, 1:])
  66         return (Q * R).sum(-1), R / ctx.A_star, U
  67
  68
  69 pscan = PScan.apply
  70
  71 ######################################################################
  72
  73 if __name__ == "__main__":
  74     A = torch.randn(1, 5, dtype=torch.float64).requires_grad_()
  75     X = torch.randn(1, 5, 3, dtype=torch.float64).requires_grad_()
  76     Y0 = torch.randn(1, 3, dtype=torch.float64).requires_grad_()
  77
  78     y = Y0[:, None]
  79
  80     for k in range(A.size(1)):
  81         y = A[:, k, None] * y + X[:, k]
  82         print(f"{k} -> {y}")
  83
  84     print(torch.autograd.grad(y.mean(), A, retain_graph=True))
  85     print(torch.autograd.grad(y.mean(), X, retain_graph=True))
  86     print(torch.autograd.grad(y.mean(), Y0, retain_graph=True))
  87
  88     Y = pscan(A, X, Y0)
  89
  90     print()
  91
  92     for k in range(A.size(1)):
  93         print(f"{k} -> {Y[:,k]}")
  94
  95     y = Y[:, -1]
  96
  97     print(torch.autograd.grad(y.mean(), A, retain_graph=True))
  98     print(torch.autograd.grad(y.mean(), X, retain_graph=True))
  99     print(torch.autograd.grad(y.mean(), Y0, retain_graph=True))