Skip to content

Commit 2b655cf

Browse files
committed
add a fast-path for size == 1 arrays in sort
radix sort implementation asserts that array must be of size > 1
1 parent f7e0967 commit 2b655cf

File tree

1 file changed

+2
-9
lines changed

1 file changed

+2
-9
lines changed

dpctl/tensor/_sorting.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ def sort(x, /, *, axis=-1, descending=False, stable=True, kind=None):
8585
return dpt.copy(x, order="C")
8686
else:
8787
axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis")
88+
if x.size == 1:
89+
return dpt.copy(x, order="C")
8890
a1 = axis + 1
8991
if a1 == nd:
9092
perm = list(range(nd))
@@ -96,15 +98,6 @@ def sort(x, /, *, axis=-1, descending=False, stable=True, kind=None):
9698
arr = dpt.permute_dims(x, perm)
9799
if kind is None:
98100
kind = "stable"
99-
if not isinstance(kind, str) or kind not in [
100-
"stable",
101-
"radixsort",
102-
"mergesort",
103-
]:
104-
raise ValueError(
105-
"Unsupported kind value. Expected 'stable', 'mergesort', "
106-
f"or 'radixsort', but got '{kind}'"
107-
)
108101
if kind == "mergesort":
109102
impl_fn = _get_mergesort_impl_fn(descending)
110103
elif kind == "radixsort":

0 commit comments

Comments
 (0)