add rtol kwarg to dpnp.linalg.pinv and dpnp.linalg.matrix_rank (#2124)

* add rtol keyword argument to dpnp.linalg.pinv and dpnp.linalg.matrix_rank

* address comments

Author: vtavana

dpnp/linalg/dpnp_iface_linalg.py

@@ -99,7 +99,7 @@
 ]
 
 
-def cholesky(a, upper=False):
+def cholesky(a, /, *, upper=False):
     """
     Cholesky decomposition.
 
@@ -1062,7 +1062,7 @@ def matrix_power(a, n):
     return dpnp_matrix_power(a, n)
 
 
-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.
 
@@ -1073,21 +1073,27 @@ def matrix_rank(A, tol=None, hermitian=False):
     ----------
     A : {(M,), (..., M, N)} {dpnp.ndarray, usm_ndarray}
         Input vector or stack of matrices.
-    tol : (...) {float, dpnp.ndarray, usm_ndarray}, optional
-        Threshold below which SVD values are considered zero. If `tol` is
-        ``None``, and ``S`` is an array with singular values for `M`, and
-        ``eps`` is the epsilon value for datatype of ``S``, then `tol` is
-        set to ``S.max() * max(M.shape) * eps``.
+    tol : (...) {None, float, dpnp.ndarray, usm_ndarray}, optional
+        Threshold below which SVD values are considered zero. Only `tol` or
+        `rtol` can be set at a time. If none of them are provided, defaults
+        to ``S.max() * max(M, N) * eps`` where `S` is an array with singular
+        values for `A`, and `eps` is the epsilon value for datatype of `S`.
         Default: ``None``.
     hermitian : bool, optional
         If ``True``, `A` is assumed to be Hermitian (symmetric if real-valued),
         enabling a more efficient method for finding singular values.
         Default: ``False``.
+    rtol : (...) {None, float, dpnp.ndarray, usm_ndarray}, optional
+        Parameter for the relative tolerance component. Only `tol` or `rtol`
+        can be set at a time. If none of them are provided, defaults to
+        ``max(M, N) * eps`` where `eps` is the epsilon value for datatype
+        of `S` (an array with singular values for `A`).
+        Default: ``None``.
 
     Returns
     -------
     rank : (...) dpnp.ndarray
-        Rank of A.
+        Rank of `A`.
 
     See Also
     --------
@@ -1114,8 +1120,12 @@ def matrix_rank(A, tol=None, hermitian=False):
         dpnp.check_supported_arrays_type(
             tol, scalar_type=True, all_scalars=True
         )
+    if rtol is not None:
+        dpnp.check_supported_arrays_type(
+            rtol, scalar_type=True, all_scalars=True
+        )
 
-    return dpnp_matrix_rank(A, tol=tol, hermitian=hermitian)
+    return dpnp_matrix_rank(A, tol=tol, hermitian=hermitian, rtol=rtol)
 
 
 def matrix_transpose(x, /):
@@ -1456,7 +1466,7 @@ def outer(x1, x2, /):
     return dpnp.outer(x1, x2)
 
 
-def pinv(a, rcond=1e-15, hermitian=False):
+def pinv(a, rcond=None, hermitian=False, *, rtol=None):
     """
     Compute the (Moore-Penrose) pseudo-inverse of a matrix.
 
@@ -1469,20 +1479,27 @@ def pinv(a, rcond=1e-15, hermitian=False):
     ----------
     a : (..., M, N) {dpnp.ndarray, usm_ndarray}
         Matrix or stack of matrices to be pseudo-inverted.
-    rcond : {float, dpnp.ndarray, usm_ndarray}, optional
+    rcond : (...) {None, float, dpnp.ndarray, usm_ndarray}, optional
         Cutoff for small singular values.
         Singular values less than or equal to ``rcond * largest_singular_value``
         are set to zero. Broadcasts against the stack of matrices.
-        Default: ``1e-15``.
+        Only `rcond` or `rtol` can be set at a time. If none of them are
+        provided, defaults to ``max(M, N) * dpnp.finfo(a.dtype).eps``.
+        Default: ``None``.
     hermitian : bool, optional
         If ``True``, a is assumed to be Hermitian (symmetric if real-valued),
         enabling a more efficient method for finding singular values.
         Default: ``False``.
+    rtol : (...) {None, float, dpnp.ndarray, usm_ndarray}, optional
+        Same as `rcond`, but it's an Array API compatible parameter name.
+        Only `rcond` or `rtol` can be set at a time. If none of them are
+        provided, defaults to ``max(M, N) * dpnp.finfo(a.dtype).eps``.
+        Default: ``None``.
 
     Returns
     -------
     out : (..., N, M) dpnp.ndarray
-        The pseudo-inverse of a.
+        The pseudo-inverse of `a`.
 
     Examples
     --------
@@ -1493,17 +1510,24 @@ def pinv(a, rcond=1e-15, hermitian=False):
     >>> a = np.random.randn(9, 6)
     >>> B = np.linalg.pinv(a)
     >>> np.allclose(a, np.dot(a, np.dot(B, a)))
-    array([ True])
+    array(True)
     >>> np.allclose(B, np.dot(B, np.dot(a, B)))
-    array([ True])
+    array(True)
 
     """
 
     dpnp.check_supported_arrays_type(a)
-    dpnp.check_supported_arrays_type(rcond, scalar_type=True, all_scalars=True)
+    if rcond is not None:
+        dpnp.check_supported_arrays_type(
+            rcond, scalar_type=True, all_scalars=True
+        )
+    if rtol is not None:
+        dpnp.check_supported_arrays_type(
+            rtol, scalar_type=True, all_scalars=True
+        )
     assert_stacked_2d(a)
 
-    return dpnp_pinv(a, rcond=rcond, hermitian=hermitian)
+    return dpnp_pinv(a, rcond=rcond, hermitian=hermitian, rtol=rtol)
 
 
 def qr(a, mode="reduced"):

dpnp/linalg/dpnp_utils_linalg.py

@@ -2214,24 +2214,33 @@ def dpnp_matrix_power(a, n):
     return result
 
 
-def dpnp_matrix_rank(A, tol=None, hermitian=False):
+def dpnp_matrix_rank(A, tol=None, hermitian=False, rtol=None):
     """
-    dpnp_matrix_rank(A, tol=None, hermitian=False)
+    dpnp_matrix_rank(A, tol=None, hermitian=False, rtol=None)
 
     Return matrix rank of array using SVD method.
 
     """
 
+    if rtol is not None and tol is not None:
+        raise ValueError("`tol` and `rtol` can't be both set.")
+
     if A.ndim < 2:
         return (A != 0).any().astype(int)
 
     S = dpnp_svd(A, compute_uv=False, hermitian=hermitian)
 
     if tol is None:
-        rtol = max(A.shape[-2:]) * dpnp.finfo(S.dtype).eps
+        if rtol is None:
+            rtol = max(A.shape[-2:]) * dpnp.finfo(S.dtype).eps
+        elif not dpnp.isscalar(rtol):
+            # Add a new axis to make it broadcastable against S
+            # needed for S > tol comparison below
+            rtol = rtol[..., None]
         tol = S.max(axis=-1, keepdims=True) * rtol
     elif not dpnp.isscalar(tol):
-        # Add a new axis to match NumPy's output
+        # Add a new axis to make it broadcastable against S,
+        # needed for S > tol comparison below
         tol = tol[..., None]
 
     return dpnp.count_nonzero(S > tol, axis=-1)
@@ -2320,9 +2329,9 @@ def dpnp_norm(x, ord=None, axis=None, keepdims=False):
     raise ValueError("Improper number of dimensions to norm.")
 
 
-def dpnp_pinv(a, rcond=1e-15, hermitian=False):
+def dpnp_pinv(a, rcond=None, hermitian=False, rtol=None):
     """
-    dpnp_pinv(a, rcond=1e-15, hermitian=False):
+    dpnp_pinv(a, rcond=None, hermitian=False, rtol=None)
 
     Compute the Moore-Penrose pseudoinverse of `a` matrix.
 
@@ -2331,6 +2340,15 @@ def dpnp_pinv(a, rcond=1e-15, hermitian=False):
 
     """
 
+    if rcond is None:
+        if rtol is None:
+            dtype = dpnp.result_type(a.dtype, dpnp.default_float_type(a.device))
+            rcond = max(a.shape[-2:]) * dpnp.finfo(dtype).eps
+        else:
+            rcond = rtol
+    elif rtol is not None:
+        raise ValueError("`rtol` and `rcond` can't be both set.")
+
     if _is_empty_2d(a):
         m, n = a.shape[-2:]
         sh = a.shape[:-2] + (n, m)

tests/test_linalg.py

@@ -2097,6 +2097,32 @@ def test_matrix_rank_tolerance(self, high_tol, low_tol):
         )
         assert np_rank_low_tol == dp_rank_low_tol
 
+    # rtol kwarg was added in numpy 2.0
+    @testing.with_requires("numpy>=2.0")
+    @pytest.mark.parametrize(
+        "tol",
+        [0.99e-6, numpy.array(1.01e-6), numpy.ones(4) * [0.99e-6]],
+        ids=["float", "0-D array", "1-D array"],
+    )
+    def test_matrix_rank_tol(self, tol):
+        a = numpy.zeros((4, 3, 2))
+        a_dp = inp.array(a)
+
+        if isinstance(tol, numpy.ndarray):
+            dp_tol = inp.array(
+                tol, usm_type=a_dp.usm_type, sycl_queue=a_dp.sycl_queue
+            )
+        else:
+            dp_tol = tol
+
+        expected = numpy.linalg.matrix_rank(a, rtol=tol)
+        result = inp.linalg.matrix_rank(a_dp, rtol=dp_tol)
+        assert_dtype_allclose(result, expected)
+
+        expected = numpy.linalg.matrix_rank(a, tol=tol)
+        result = inp.linalg.matrix_rank(a_dp, tol=dp_tol)
+        assert_dtype_allclose(result, expected)
+
     def test_matrix_rank_errors(self):
         a_dp = inp.array([[1, 2], [3, 4]], dtype="float32")
 
@@ -2121,6 +2147,11 @@ def test_matrix_rank_errors(self):
             tol_dp_q,
         )
 
+        # both tol and rtol are given
+        assert_raises(
+            ValueError, inp.linalg.matrix_rank, a_dp, tol=1e-06, rtol=1e-04
+        )
+
 
 # numpy.linalg.matrix_transpose() is available since numpy >= 2.0
 @testing.with_requires("numpy>=2.0")
@@ -3141,6 +3172,16 @@ def test_pinv_hermitian(self, dtype, shape):
         reconstructed = inp.dot(inp.dot(a_dp, B_dp), a_dp)
         assert_allclose(reconstructed, a_dp, rtol=tol, atol=tol)
 
+    # rtol kwarg was added in numpy 2.0
+    @testing.with_requires("numpy>=2.0")
+    def test_pinv_rtol(self):
+        a = numpy.ones((2, 2))
+        a_dp = inp.array(a)
+
+        expected = numpy.linalg.pinv(a, rtol=1e-15)
+        result = inp.linalg.pinv(a_dp, rtol=1e-15)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
     @pytest.mark.parametrize(
         "shape",
@@ -3207,6 +3248,11 @@ def test_pinv_errors(self):
         rcond_dp_q = inp.array([0.5], dtype="float32", sycl_queue=rcond_queue)
         assert_raises(ValueError, inp.linalg.pinv, a_dp_q, rcond_dp_q)
 
+        # both rcond and rtol are given
+        assert_raises(
+            ValueError, inp.linalg.pinv, a_dp, rcond=1e-06, rtol=1e-04
+        )
+
 
 class TestTensorinv:
     @pytest.mark.parametrize("dtype", get_all_dtypes())