pscan.py

   1 #!/usr/bin/env python
   2
   3 import torch
   4
   5 ######################################################################
   6
   7
   8 def naive_rec(A, X, Y0):
   9     Y = []
  10     for t in range(X.size(1)):
  11         if t == 0:
  12             Y.append(A[:, t] * Y0 + X[:, t])
  13         else:
  14             Y.append(A[:, t] * Y[-1] + X[:, t])
  15
  16     return torch.cat([y[:, None, :] for y in Y], dim=1)
  17
  18
  19 ######################################################################
  20
  21 # A is NxTx1
  22 # X is NxTxD
  23 # Y0 is NxD
  24 #
  25 # Returns Y defined with
  26 #
  27 #           Y[:, 0] = A[:, 0] * Y0 + X[:,0]
  28 # for t > 0 Y[:, t] = A[:, t] * Y[:, t - 1] + X[:, t]
  29
  30
  31 def pscan_rec(A, X, Y0):
  32     if X.size(1) % 2 == 1:
  33         if X.size(1) == 1:
  34             return A[:, :1] * Y0[:, None] + X[:, :1]
  35         else:
  36             Y = pscan_rec(A[:, :-1], X[:, :-1], Y0)
  37             return torch.cat([Y, A[:, -1:] * Y[:, -1:] + X[:, -1:]], dim=1)
  38
  39     A2 = A.reshape(A.size(0), A.size(1) // 2, 2, A.size(2))
  40     X2 = X.reshape(X.size(0), X.size(1) // 2, 2, X.size(2))
  41
  42     X_star = X2[:, :, 0].clone()
  43     X_star[:, 1:] += A2[:, 1:, 0] * X2[:, :-1, 1]
  44
  45     A_star = A2[:, :, 0].clone()
  46     A_star[:, 1:] *= A2[:, :-1, 1]
  47
  48     Y_star = pscan_rec(A_star, X_star, Y0)[:, :, None]
  49
  50     Y = torch.cat([Y_star, A2[:, :, 1, None] * Y_star + X2[:, :, 1, None]], dim=2)
  51
  52     Y = Y.reshape(Y.size(0), -1, Y.size(-1))
  53
  54     return Y
  55
  56
  57 ######################################################################
  58
  59 N, T, D = 5, 29, 12
  60
  61 A = torch.rand(N, T, 1, dtype=torch.float64)
  62 X = torch.randint(10, (N, T, D), dtype=torch.float64)
  63 Y0 = torch.randint(10, (N, D), dtype=torch.float64)
  64
  65 naive_Y = naive_rec(A, X, Y0)
  66
  67 pscan_Y = pscan_rec(A, X, Y0)
  68
  69 print((naive_Y - pscan_Y).pow(2).mean())
  70
  71 pscan_Y1 = pscan_rec(A[:, :15], X[:, :15], Y0)
  72 pscan_Y2 = pscan_rec(A[:, 15:], X[:, 15:], pscan_Y1[:, -1])
  73
  74 print((naive_Y - torch.cat([pscan_Y1, pscan_Y2], dim=1)).pow(2).mean())