numpy · ngoldbaum · Jan 18, 2024 · Dec 21, 2023 · Jan 7, 2024 · Jan 17, 2024
diff --git a/doc/release/upcoming_changes/25437.new_feature.rst b/doc/release/upcoming_changes/25437.new_feature.rst
@@ -0,0 +1,8 @@
+`numpy.linalg.martrix_rank`, `numpy.sort` and `numpy.argsort` new parameters
+----------------------------------------------------------------------------
+
+New keyword parameters were added to improve array API compatibility:
+
+* ``rtol`` keyword parameter was added to `numpy.linalg.martrix_rank`.
+
+* ``stable`` keyword parameter was added to `numpy.sort` and `numpy.argsort`.
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -1113,6 +1113,8 @@ class _ArrayOrScalarCommon:
         axis: None | SupportsIndex = ...,
         kind: None | _SortKind = ...,
         order: None | str | Sequence[str] = ...,
+        *,
+        stable: None | bool = ...,
     ) -> NDArray[Any]: ...
 
     @overload
@@ -1640,6 +1642,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
         axis: SupportsIndex = ...,
         kind: None | _SortKind = ...,
         order: None | str | Sequence[str] = ...,
+        *,
+        stable: None | bool = ...,
     ) -> None: ...
 
     @overload

diff --git a/numpy/_core/fromnumeric.py b/numpy/_core/fromnumeric.py
@@ -894,12 +894,12 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
     return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order)
 
 
-def _sort_dispatcher(a, axis=None, kind=None, order=None):
+def _sort_dispatcher(a, axis=None, kind=None, order=None, *, stable=None):
     return (a,)
 
 
 @array_function_dispatch(_sort_dispatcher)
-def sort(a, axis=-1, kind=None, order=None):
+def sort(a, axis=-1, kind=None, order=None, *, stable=None):
     """
     Return a sorted copy of an array.
 
@@ -925,6 +925,13 @@ def sort(a, axis=-1, kind=None, order=None):
         be specified as a string, and not all fields need be specified,
         but unspecified fields will still be used, in the order in which
         they come up in the dtype, to break ties.
+    stable : bool, optional
+        Sort stability. If ``True``, the returned array will maintain
+        the relative order of ``a`` values which compare as equal.
+        If ``False`` or ``None``, this is not guaranteed. Internally,
+        this option selects ``kind='stable'``. Default: ``None``.
+
+        .. versionadded:: 2.0.0
 
     Returns
     -------
@@ -1053,16 +1060,16 @@ def sort(a, axis=-1, kind=None, order=None):
         axis = -1
     else:
         a = asanyarray(a).copy(order="K")
-    a.sort(axis=axis, kind=kind, order=order)
+    a.sort(axis=axis, kind=kind, order=order, stable=stable)
     return a
 
 
-def _argsort_dispatcher(a, axis=None, kind=None, order=None):
+def _argsort_dispatcher(a, axis=None, kind=None, order=None, *, stable=None):
     return (a,)
 
 
 @array_function_dispatch(_argsort_dispatcher)
-def argsort(a, axis=-1, kind=None, order=None):
+def argsort(a, axis=-1, kind=None, order=None, *, stable=None):
     """
     Returns the indices that would sort an array.
 
@@ -1091,6 +1098,13 @@ def argsort(a, axis=-1, kind=None, order=None):
         be specified as a string, and not all fields need be specified,
         but unspecified fields will still be used, in the order in which
         they come up in the dtype, to break ties.
+    stable : bool, optional
+        Sort stability. If ``True``, the returned array will maintain
+        the relative order of ``a`` values which compare as equal.
+        If ``False`` or ``None``, this is not guaranteed. Internally,
+        this option selects ``kind='stable'``. Default: ``None``.
+
+        .. versionadded:: 2.0.0
 
     Returns
     -------
@@ -1169,8 +1183,9 @@ def argsort(a, axis=-1, kind=None, order=None):
     array([0, 1])
 
     """
-    return _wrapfunc(a, 'argsort', axis=axis, kind=kind, order=order)
-
+    return _wrapfunc(
+        a, 'argsort', axis=axis, kind=kind, order=order, stable=stable
+    )
 
 def _argmax_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue):
     return (a, out)

diff --git a/numpy/_core/fromnumeric.pyi b/numpy/_core/fromnumeric.pyi
@@ -211,20 +211,26 @@ def sort(
     axis: None | SupportsIndex = ...,
     kind: None | _SortKind = ...,
     order: None | str | Sequence[str] = ...,
+    *,
+    stable: None | bool = ...,
 ) -> NDArray[_SCT]: ...
 @overload
 def sort(
     a: ArrayLike,
     axis: None | SupportsIndex = ...,
     kind: None | _SortKind = ...,
     order: None | str | Sequence[str] = ...,
+    *,
+    stable: None | bool = ...,
 ) -> NDArray[Any]: ...
 
 def argsort(
     a: ArrayLike,
     axis: None | SupportsIndex = ...,
     kind: None | _SortKind = ...,
     order: None | str | Sequence[str] = ...,
+    *,
+    stable: None | bool = ...,
 ) -> NDArray[intp]: ...
 
 @overload

diff --git a/numpy/_core/include/numpy/ndarraytypes.h b/numpy/_core/include/numpy/ndarraytypes.h
@@ -155,6 +155,7 @@ enum NPY_TYPECHAR {
  * depend on the data type.
  */
 typedef enum {
+        _NPY_SORT_UNDEFINED=-1,
         NPY_QUICKSORT=0,
         NPY_HEAPSORT=1,
         NPY_MERGESORT=2,

diff --git a/numpy/_core/src/multiarray/_multiarray_tests.c.src b/numpy/_core/src/multiarray/_multiarray_tests.c.src
@@ -2032,6 +2032,7 @@ run_sortkind_converter(PyObject* NPY_UNUSED(self), PyObject *args)
         return NULL;
     }
     switch (kind) {
+        case _NPY_SORT_UNDEFINED: return PyUnicode_FromString("_NPY_SORT_UNDEFINED");
         case NPY_QUICKSORT: return PyUnicode_FromString("NPY_QUICKSORT");
         case NPY_HEAPSORT: return PyUnicode_FromString("NPY_HEAPSORT");
         case NPY_STABLESORT: return PyUnicode_FromString("NPY_STABLESORT");

diff --git a/numpy/_core/src/multiarray/conversion_utils.c b/numpy/_core/src/multiarray/conversion_utils.c
@@ -431,6 +431,28 @@ PyArray_BoolConverter(PyObject *object, npy_bool *val)
     return NPY_SUCCEED;
 }
 
+/*
+ * Optionally convert an object to true / false
+ */
+NPY_NO_EXPORT int
+PyArray_OptionalBoolConverter(PyObject *object, int *val)
+{
+    /* Leave the desired default from the caller for Py_None */
+    if (object == Py_None) {
+        return NPY_SUCCEED;
+    }
+    if (PyObject_IsTrue(object)) {
+        *val = 1;
+    }
+    else {
+        *val = 0;
+    }
+    if (PyErr_Occurred()) {
+        return NPY_FAIL;
+    }
+    return NPY_SUCCEED;
+}
+
 static int
 string_converter_helper(
     PyObject *object,

diff --git a/numpy/_core/src/multiarray/conversion_utils.h b/numpy/_core/src/multiarray/conversion_utils.h
@@ -27,6 +27,9 @@ PyArray_BufferConverter(PyObject *obj, PyArray_Chunk *buf);
 NPY_NO_EXPORT int
 PyArray_BoolConverter(PyObject *object, npy_bool *val);
 
+NPY_NO_EXPORT int
+PyArray_OptionalBoolConverter(PyObject *object, int *val);
+
 NPY_NO_EXPORT int
 PyArray_ByteorderConverter(PyObject *obj, char *endian);
 

diff --git a/numpy/_core/src/multiarray/methods.c b/numpy/_core/src/multiarray/methods.c
@@ -1235,18 +1235,20 @@ static PyObject *
 array_sort(PyArrayObject *self,
         PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
 {
-    int axis=-1;
+    int axis = -1;
     int val;
-    NPY_SORTKIND sortkind = NPY_QUICKSORT;
+    NPY_SORTKIND sortkind = _NPY_SORT_UNDEFINED;
     PyObject *order = NULL;
     PyArray_Descr *saved = NULL;
     PyArray_Descr *newd;
+    int stable = -1;
     NPY_PREPARE_ARGPARSER;
 
     if (npy_parse_arguments("sort", args, len_args, kwnames,
             "|axis", &PyArray_PythonPyIntFromInt, &axis,
             "|kind", &PyArray_SortkindConverter, &sortkind,
             "|order", NULL, &order,
+            "$stable", &PyArray_OptionalBoolConverter, &stable,
             NULL, NULL, NULL) < 0) {
         return NULL;
     }
@@ -1281,6 +1283,18 @@ array_sort(PyArrayObject *self,
         newd->names = new_name;
         ((PyArrayObject_fields *)self)->descr = newd;
     }
+    if (sortkind != _NPY_SORT_UNDEFINED && stable != -1) {
+        PyErr_SetString(PyExc_ValueError,
+            "`kind` and `stable` parameters can't be provided at "
+            "the same time. Use only one of them.");
+        return NULL;
+    }
+    else if ((sortkind == _NPY_SORT_UNDEFINED && stable == -1) || (stable == 0)) {
+        sortkind = NPY_QUICKSORT;
+    }
+    else if (stable == 1) {
+        sortkind = NPY_STABLESORT;
+    }
 
     val = PyArray_Sort(self, axis, sortkind);
     if (order != NULL) {
@@ -1371,15 +1385,17 @@ array_argsort(PyArrayObject *self,
         PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
 {
     int axis = -1;
-    NPY_SORTKIND sortkind = NPY_QUICKSORT;
+    NPY_SORTKIND sortkind = _NPY_SORT_UNDEFINED;
     PyObject *order = NULL, *res;
     PyArray_Descr *newd, *saved=NULL;
+    int stable = -1;
     NPY_PREPARE_ARGPARSER;
 
     if (npy_parse_arguments("argsort", args, len_args, kwnames,
             "|axis", &PyArray_AxisConverter, &axis,
             "|kind", &PyArray_SortkindConverter, &sortkind,
             "|order", NULL, &order,
+            "$stable", &PyArray_OptionalBoolConverter, &stable,
             NULL, NULL, NULL) < 0) {
         return NULL;
     }
@@ -1414,6 +1430,18 @@ array_argsort(PyArrayObject *self,
         newd->names = new_name;
         ((PyArrayObject_fields *)self)->descr = newd;
     }
+    if (sortkind != _NPY_SORT_UNDEFINED && stable != -1) {
+        PyErr_SetString(PyExc_ValueError,
+            "`kind` and `stable` parameters can't be provided at "
+            "the same time. Use only one of them.");
+        return NULL;
+    }
+    else if ((sortkind == _NPY_SORT_UNDEFINED && stable == -1) || (stable == 0)) {
+        sortkind = NPY_QUICKSORT;
+    }
+    else if (stable == 1) {
+        sortkind = NPY_STABLESORT;
+    }
 
     res = PyArray_ArgSort(self, axis, sortkind);
     if (order != NULL) {

diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
@@ -2043,6 +2043,12 @@ def test_sort(self):
         b = np.sort(a)
         assert_equal(b, a[::-1], msg)
 
+        with assert_raises_regex(
+            ValueError,
+            "kind` and `stable` parameters can't be provided at the same time"
+        ):
+            np.sort(a, kind="stable", stable=True)
+
     # all c scalar sorts use the same code with different types
     # so it suffices to run a quick check with one type. The number
     # of sorted items must be greater than ~50 to check the actual
@@ -2481,6 +2487,12 @@ def test_argsort(self):
         a = np.array(['aaaaaaaaa' for i in range(100)], dtype=np.str_)
         assert_equal(a.argsort(kind='m'), r)
 
+        with assert_raises_regex(
+            ValueError,
+            "kind` and `stable` parameters can't be provided at the same time"
+        ):
+            np.argsort(a, kind="stable", stable=True)
+
     def test_sort_unicode_kind(self):
         d = np.arange(10)
         k = b'\xc3\xa4'.decode("UTF8")

diff --git a/numpy/linalg/_linalg.py b/numpy/linalg/_linalg.py
@@ -1965,12 +1965,12 @@ def cond(x, p=None):
     return r
 
 
-def _matrix_rank_dispatcher(A, tol=None, hermitian=None):
+def _matrix_rank_dispatcher(A, tol=None, hermitian=None, *, rtol=None):
     return (A,)
 
 
 @array_function_dispatch(_matrix_rank_dispatcher)
-def matrix_rank(A, tol=None, hermitian=False):
+def matrix_rank(A, tol=None, hermitian=False, *, rtol=None):
     """
     Return matrix rank of array using SVD method
 
@@ -1998,6 +1998,11 @@ def matrix_rank(A, tol=None, hermitian=False):
         Defaults to False.
 
         .. versionadded:: 1.14
+    rtol : (...) array_like, float, optional
+        Parameter for the relative tolerance component. Only ``tol`` or
+        ``rtol`` can be set at a time. Defaults to ``max(M, N) * eps``.
+
+        .. versionadded:: 2.0.0
 
     Returns
     -------
@@ -2067,12 +2072,12 @@ def matrix_rank(A, tol=None, hermitian=False):
     if A.ndim < 2:
         return int(not all(A == 0))
     S = svd(A, compute_uv=False, hermitian=hermitian)
+    if rtol is not None and tol is not None:
+        raise ValueError("`tol` and `rtol` can't be both set.")
+    if rtol is None:
+        rtol = max(A.shape[-2:]) * finfo(S.dtype).eps
     if tol is None:
-        tol = (
-            S.max(axis=-1, keepdims=True) *
-            max(A.shape[-2:]) *
-            finfo(S.dtype).eps
-        )
+        tol = S.max(axis=-1, keepdims=True) * rtol
     else:
         tol = asarray(tol)[..., newaxis]
     return count_nonzero(S > tol, axis=-1)

diff --git a/numpy/linalg/_linalg.pyi b/numpy/linalg/_linalg.pyi
@@ -265,6 +265,8 @@ def matrix_rank(
     A: _ArrayLikeComplex_co,
     tol: None | _ArrayLikeFloat_co = ...,
     hermitian: bool = ...,
+    *,
+    rtol: None | _ArrayLikeFloat_co = ...,
 ) -> Any: ...
 
 @overload

diff --git a/numpy/linalg/tests/test_linalg.py b/numpy/linalg/tests/test_linalg.py
@@ -1624,6 +1624,11 @@ def test_matrix_rank(self):
         # works on scalar
         assert_equal(matrix_rank(1), 1)
 
+        with assert_raises_regex(
+            ValueError, "`tol` and `rtol` can\'t be both set."
+        ):
+            matrix_rank(I, tol=0.01, rtol=0.01)
+
     def test_symmetric_rank(self):
         assert_equal(4, matrix_rank(np.eye(4), hermitian=True))
         assert_equal(1, matrix_rank(np.ones((4, 4)), hermitian=True))