Skip to content

Commit 1084835

Browse files
committed
comment about peak memory
1 parent 2a3fd9f commit 1084835

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

src/array_api_extra/_delegation.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -714,7 +714,7 @@ def partition(
714714
ranks = xp.arange(a.shape[-1]).expand_as(a)
715715

716716
split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True)
717-
del indices
717+
del indices # indices won't be used => del ASAP to reduce peak memory usage
718718

719719
# fill the left-side of the partition
720720
mask_src = a < split_value
@@ -811,7 +811,7 @@ def argpartition(
811811
out = xp.empty_like(ranks)
812812

813813
split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True)
814-
del indices
814+
del indices # indices won't be used => del ASAP to reduce peak memory usage
815815

816816
mask_src = a < split_value
817817
n_left = mask_src.sum(dim=-1, keepdim=True)

0 commit comments

Comments
 (0)