Array API compatible wrapper for :py:func:`np.argsort `. See its docstring for more information.
(
x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
)
| 8 | |
| 9 | # Note: the descending keyword argument is new in this function |
| 10 | def argsort( |
| 11 | x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True |
| 12 | ) -> Array: |
| 13 | """ |
| 14 | Array API compatible wrapper for :py:func:`np.argsort <numpy.argsort>`. |
| 15 | |
| 16 | See its docstring for more information. |
| 17 | """ |
| 18 | if x.dtype not in _real_numeric_dtypes: |
| 19 | raise TypeError("Only real numeric dtypes are allowed in argsort") |
| 20 | # Note: this keyword argument is different, and the default is different. |
| 21 | kind = "stable" if stable else "quicksort" |
| 22 | if not descending: |
| 23 | res = np.argsort(x._array, axis=axis, kind=kind) |
| 24 | else: |
| 25 | # As NumPy has no native descending sort, we imitate it here. Note that |
| 26 | # simply flipping the results of np.argsort(x._array, ...) would not |
| 27 | # respect the relative order like it would in native descending sorts. |
| 28 | res = np.flip( |
| 29 | np.argsort(np.flip(x._array, axis=axis), axis=axis, kind=kind), |
| 30 | axis=axis, |
| 31 | ) |
| 32 | # Rely on flip()/argsort() to validate axis |
| 33 | normalised_axis = axis if axis >= 0 else x.ndim + axis |
| 34 | max_i = x.shape[normalised_axis] - 1 |
| 35 | res = max_i - res |
| 36 | return Array._new(res) |
| 37 | |
| 38 | # Note: the descending keyword argument is new in this function |
| 39 | def sort( |