We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 2a3fd9f commit 1084835Copy full SHA for 1084835
src/array_api_extra/_delegation.py
@@ -714,7 +714,7 @@ def partition(
714
ranks = xp.arange(a.shape[-1]).expand_as(a)
715
716
split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True)
717
- del indices
+ del indices # indices won't be used => del ASAP to reduce peak memory usage
718
719
# fill the left-side of the partition
720
mask_src = a < split_value
@@ -811,7 +811,7 @@ def argpartition(
811
out = xp.empty_like(ranks)
812
813
814
815
816
817
n_left = mask_src.sum(dim=-1, keepdim=True)
0 commit comments