fixes

unixpickle · unixpickle · commit 17baeabb1734 · 2024-03-17T22:31:16.000-04:00
diff --git a/accelerated_scan/warp.cuh b/accelerated_scan/warp.cuh
@@ -356,6 +356,8 @@ __device__ Tuple load_shifted_tuple(const Tuple* ptr, int index, int limit) {
         const int idx = index * Tuple::Size + i + offset;
         if (idx >= 0 && idx < limit * Tuple::Size) {
             x.data[i] = rawPtr[offset];
+        } else {
+            x.data[i] = 0;
         }
     }
 
diff --git a/accelerated_scan/warp.py b/accelerated_scan/warp.py
@@ -7,6 +7,7 @@
 
 cpp_source = """
 at::Tensor warpscan_forward(const at::Tensor &gates, const at::Tensor &tokens, const at::Tensor &out, const bool reverse);
+void warpscan_backward(const at::Tensor &gates, const at::Tensor &output, const at::Tensor &outGrad, const at::Tensor& gateGradOut, const at::Tensor& valueGradOut);
 """
 
 module = load_inline(

Original file line number	Diff line number	Diff line change
`@@ -356,6 +356,8 @@ __device__ Tuple load_shifted_tuple(const Tuple* ptr, int index, int limit) {`
`356`	`356`	`const int idx = index * Tuple::Size + i + offset;`
`357`	`357`	`if (idx >= 0 && idx < limit * Tuple::Size) {`
`358`	`358`	`x.data[i] = rawPtr[offset];`
	`359`	`+ } else {`
	`360`	`+ x.data[i] = 0;`
`359`	`361`	`}`
`360`	`362`	`}`
`361`	`363`