Avoid gpu searchsorted when k==1

TDAmeritrade · Dec 1, 2022 · 919ec76 · 919ec76
1 parent 5d8a637
commit 919ec76
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/stumpy/gpu_stump.py b/stumpy/gpu_stump.py
@@ -194,13 +194,17 @@ def _compute_and_update_PI_kernel(
                 indices_R[j] = i
 
         if p_norm < profile[j, -1]:
-            idx = core._gpu_searchsorted_right(profile[j], p_norm, bfs, nlevel)
-            for g in range(k - 1, idx, -1):
-                profile[j, g] = profile[j, g - 1]
-                indices[j, g] = indices[j, g - 1]
+            if k == 1:
+                profile[j, 0] = p_norm
+                indices[j, 0] = i
+            else:
+                idx = core._gpu_searchsorted_right(profile[j], p_norm, bfs, nlevel)
+                for g in range(k - 1, idx, -1):
+                    profile[j, g] = profile[j, g - 1]
+                    indices[j, g] = indices[j, g - 1]
 
-            profile[j, idx] = p_norm
-            indices[j, idx] = i
+                profile[j, idx] = p_norm
+                indices[j, idx] = i
 
 
 def _gpu_stump(