use kernel

Author: unixpickle

accelerated_scan/warp.py

@@ -27,6 +27,7 @@
     ]
 )
 warpscan_forward = module.warpscan_forward
+warpscan_backward = module.warpscan_backward
 
 def scan_forward(gates, tokens, reverse=False):
     output = torch.zeros_like(tokens)
@@ -57,13 +58,10 @@ def backward(ctx, grad_output):
         assert states.is_contiguous()
         assert gates.is_contiguous()
 
-        padded_shifted_gates = torch.cat([gates, torch.ones_like(gates[:, :, :1])], dim=-1)[:, :, 1:].contiguous()
-        d_states = scan_forward(padded_shifted_gates, grad_output, reverse=True)
+        d_gates = torch.empty_like(gates)
+        d_tokens = torch.empty_like(gates)
+        warpscan_backward(gates, states, grad_output, d_gates, d_tokens)
 
-        padded_outputs = torch.cat([torch.zeros_like(states[:, :, :1]), states], dim=-1)[:, :, :-1]
-        d_gates = padded_outputs * d_states
-
-        d_tokens = d_states
         return d_gates, d_tokens