diff --git a/dpnp/linalg/dpnp_iface_linalg.py b/dpnp/linalg/dpnp_iface_linalg.py index 2c09bb40ed3..ae98c395641 100644 --- a/dpnp/linalg/dpnp_iface_linalg.py +++ b/dpnp/linalg/dpnp_iface_linalg.py @@ -1073,7 +1073,7 @@ def matrix_rank(A, tol=None, hermitian=False, *, rtol=None): ---------- A : {(M,), (..., M, N)} {dpnp.ndarray, usm_ndarray} Input vector or stack of matrices. - tol : (...) {float, dpnp.ndarray, usm_ndarray}, optional + 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 @@ -1083,7 +1083,7 @@ def matrix_rank(A, tol=None, hermitian=False, *, rtol=None): If ``True``, `A` is assumed to be Hermitian (symmetric if real-valued), enabling a more efficient method for finding singular values. Default: ``False``. - rtol : (...) {float, dpnp.ndarray, usm_ndarray}, optional + 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 @@ -1479,7 +1479,7 @@ def pinv(a, rcond=None, hermitian=False, *, rtol=None): ---------- 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. @@ -1490,7 +1490,7 @@ def pinv(a, rcond=None, hermitian=False, *, rtol=None): If ``True``, a is assumed to be Hermitian (symmetric if real-valued), enabling a more efficient method for finding singular values. Default: ``False``. - rtol : (...) {float, dpnp.ndarray, usm_ndarray}, optional + 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``. diff --git a/dpnp/linalg/dpnp_utils_linalg.py b/dpnp/linalg/dpnp_utils_linalg.py index b022980a7b2..2d7804234de 100644 --- a/dpnp/linalg/dpnp_utils_linalg.py +++ b/dpnp/linalg/dpnp_utils_linalg.py @@ -2234,10 +2234,13 @@ def dpnp_matrix_rank(A, tol=None, hermitian=False, rtol=None): 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) diff --git a/tests/test_linalg.py b/tests/test_linalg.py index d6278c82e3f..36825d23adc 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2100,25 +2100,28 @@ def test_matrix_rank_tolerance(self, high_tol, low_tol): # rtol kwarg was added in numpy 2.0 @testing.with_requires("numpy>=2.0") @pytest.mark.parametrize( - "rtol", - [0.99e-6, numpy.array(1.01e-6), numpy.array([0.99e-6])], + "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_rtol(self, rtol): - a = numpy.eye(4) - a[-1, -1] = 1e-6 + def test_matrix_rank_tol(self, tol): + a = numpy.zeros((4, 3, 2)) a_dp = inp.array(a) - if isinstance(rtol, numpy.ndarray): - dp_rtol = inp.array( - rtol, usm_type=a_dp.usm_type, sycl_queue=a_dp.sycl_queue + if isinstance(tol, numpy.ndarray): + dp_tol = inp.array( + tol, usm_type=a_dp.usm_type, sycl_queue=a_dp.sycl_queue ) else: - dp_rtol = rtol + dp_tol = tol - expected = numpy.linalg.matrix_rank(a, rtol=rtol) - result = inp.linalg.matrix_rank(a_dp, rtol=dp_rtol) - assert expected == result + 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")