forked from numpy/numpy
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy path_sorting_functions.py
More file actions
37 lines (29 loc) · 1.07 KB
/
_sorting_functions.py
File metadata and controls
37 lines (29 loc) · 1.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from __future__ import annotations
from ._array_object import Array
import numpy as np
def argsort(
x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
) -> Array:
"""
Array API compatible wrapper for :py:func:`np.argsort <numpy.argsort>`.
See its docstring for more information.
"""
# Note: this keyword argument is different, and the default is different.
kind = "stable" if stable else "quicksort"
res = np.argsort(x._array, axis=axis, kind=kind)
if descending:
res = np.flip(res, axis=axis)
return Array._new(res)
def sort(
x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
) -> Array:
"""
Array API compatible wrapper for :py:func:`np.sort <numpy.sort>`.
See its docstring for more information.
"""
# Note: this keyword argument is different, and the default is different.
kind = "stable" if stable else "quicksort"
res = np.sort(x._array, axis=axis, kind=kind)
if descending:
res = np.flip(res, axis=axis)
return Array._new(res)