Skip to content

Commit 5de7b35

Browse files
authored
fix: don't let CuPy iterate over Index with Python for loops (#3142)
1 parent e99b641 commit 5de7b35

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

src/awkward/_slicing.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,9 @@ def _normalise_item_bool_to_int(item: Content, backend: Backend) -> Content:
498498
cumsum[1:] = item_backend.index_nplike.asarray(
499499
item_backend.nplike.cumsum(flat_mask.data)
500500
)
501-
nextoffsets = ak.index.Index(cumsum[item.offsets])
501+
502+
item_offsets = item_backend.index_nplike.asarray(item.offsets)
503+
nextoffsets = ak.index.Index(cumsum[item_offsets])
502504

503505
else:
504506
item._touch_data(recursive=False)
@@ -546,7 +548,8 @@ def _normalise_item_bool_to_int(item: Content, backend: Backend) -> Content:
546548
cumsum = item_backend.nplike.empty(expanded.shape[0] + 1, dtype=np.int64)
547549
cumsum[0] = 0
548550
cumsum[1:] = item_backend.nplike.cumsum(expanded)
549-
nextoffsets = ak.index.Index(cumsum[item.offsets])
551+
item_offsets = item_backend.index_nplike.asarray(item.offsets)
552+
nextoffsets = ak.index.Index(cumsum[item_offsets])
550553

551554
# outindex fits into the lists; non-missing are sequential
552555
outindex = ak.index.Index64(

0 commit comments

Comments
 (0)