From f954d453473b9f5b4ae66ae01a84cba2101b4d45 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Wed, 28 May 2025 14:10:34 +1000 Subject: [PATCH 01/22] wip --- sklearn/utils/tests/test_array_api.py | 28 +++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index c0322d013b90d..50f5a9e4487aa 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -588,6 +588,34 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) +from sklearn.utils import gen_even_slices +from sklearn.utils.parallel import Parallel, delayed + + +@pytest.mark.parametrize( + "array_namespace, device, dtype_name", + yield_namespace_device_dtype_combinations(), + ids=_get_namespace_device_dtype_ids, +) +def test_fill_or_add_to_diagonal_parallel(array_namespace, device, dtype_name): + """""" + xp = _array_api_for_tests(array_namespace, device) + n_samples = 10 + array_np = numpy.zeros((n_samples, n_samples), dtype=dtype_name) + array_xp = xp.asarray(array_np, device=device) + + def dumm_func(return_array, slice_): + return_array[:, slice_] = 1 + + fd = delayed(dumm_func) + with config_context(array_api_dispatch=True): + Parallel(backend="threading", n_jobs=2)( + fd(array_xp, s) for s in gen_even_slices(n_samples, 2) + ) + _fill_or_add_to_diagonal(array_xp, value=99, xp=xp, add_value=False) + print(array_xp) + + @pytest.mark.parametrize("csr_container", CSR_CONTAINERS) @pytest.mark.parametrize("dispatch", [True, False]) def test_sparse_device(csr_container, dispatch): From 230c9b5ccf254ec14b86adccba9a36df0fc5ae67 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Wed, 28 May 2025 19:16:34 +1000 Subject: [PATCH 02/22] wip --- sklearn/utils/_array_api.py | 1 + sklearn/utils/tests/test_array_api.py | 14 ++++++++------ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index a9f35516f17b6..1603ee1fc4211 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -557,6 +557,7 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): array_flat[:end:step] += value else: array_flat[:end:step] = value + print(f"{array_flat=}\n{array=}") def _is_xp_namespace(xp, name): diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index 50f5a9e4487aa..2af4ea7b840da 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -601,19 +601,21 @@ def test_fill_or_add_to_diagonal_parallel(array_namespace, device, dtype_name): """""" xp = _array_api_for_tests(array_namespace, device) n_samples = 10 - array_np = numpy.zeros((n_samples, n_samples), dtype=dtype_name) - array_xp = xp.asarray(array_np, device=device) + # ret_np = numpy.zeros((n_samples, n_samples), dtype=dtype_name) + rng = numpy.random.RandomState(0) + X_np = numpy.array(5 * rng.random_sample((n_samples, n_samples)), dtype=dtype_name) + X_xp = xp.asarray(X_np, device=device) + ret_xp = xp.ones((n_samples, n_samples), dtype=X_xp.dtype, device=device).T def dumm_func(return_array, slice_): - return_array[:, slice_] = 1 + return_array[..., slice_] = 1 fd = delayed(dumm_func) with config_context(array_api_dispatch=True): Parallel(backend="threading", n_jobs=2)( - fd(array_xp, s) for s in gen_even_slices(n_samples, 2) + fd(ret_xp, s) for s in gen_even_slices(n_samples, 2) ) - _fill_or_add_to_diagonal(array_xp, value=99, xp=xp, add_value=False) - print(array_xp) + _fill_or_add_to_diagonal(ret_xp, value=99, xp=xp, add_value=False) @pytest.mark.parametrize("csr_container", CSR_CONTAINERS) From 70d46ff9f66abd61e7369e7e815680e2b1b29407 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Wed, 28 May 2025 19:17:43 +1000 Subject: [PATCH 03/22] wip --- sklearn/utils/tests/test_array_api.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index 2af4ea7b840da..38a87418c5357 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -590,7 +590,8 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): from sklearn.utils import gen_even_slices from sklearn.utils.parallel import Parallel, delayed - +from sklearn.metrics import euclidean_distances +from sklearn.utils._array_api import _find_matching_floating_dtype @pytest.mark.parametrize( "array_namespace, device, dtype_name", @@ -601,7 +602,6 @@ def test_fill_or_add_to_diagonal_parallel(array_namespace, device, dtype_name): """""" xp = _array_api_for_tests(array_namespace, device) n_samples = 10 - # ret_np = numpy.zeros((n_samples, n_samples), dtype=dtype_name) rng = numpy.random.RandomState(0) X_np = numpy.array(5 * rng.random_sample((n_samples, n_samples)), dtype=dtype_name) X_xp = xp.asarray(X_np, device=device) From 79764889e8e9744560435a87a9d69c4ceb43b2a6 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Wed, 28 May 2025 20:14:49 +1000 Subject: [PATCH 04/22] fix fill or add diagonal --- sklearn/decomposition/_base.py | 6 ++-- sklearn/metrics/pairwise.py | 4 +-- sklearn/utils/_array_api.py | 4 ++- sklearn/utils/tests/test_array_api.py | 41 ++++++++------------------- 4 files changed, 20 insertions(+), 35 deletions(-) diff --git a/sklearn/decomposition/_base.py b/sklearn/decomposition/_base.py index 783c316b50f27..edbfd1d019cd4 100644 --- a/sklearn/decomposition/_base.py +++ b/sklearn/decomposition/_base.py @@ -47,7 +47,7 @@ def get_covariance(self): xp.asarray(0.0, device=device(exp_var), dtype=exp_var.dtype), ) cov = (components_.T * exp_var_diff) @ components_ - _fill_or_add_to_diagonal(cov, self.noise_variance_, xp) + cov = _fill_or_add_to_diagonal(cov, self.noise_variance_, xp) return cov def get_precision(self): @@ -89,10 +89,10 @@ def get_precision(self): xp.asarray(0.0, device=device(exp_var)), ) precision = components_ @ components_.T / self.noise_variance_ - _fill_or_add_to_diagonal(precision, 1.0 / exp_var_diff, xp) + precision = _fill_or_add_to_diagonal(precision, 1.0 / exp_var_diff, xp) precision = components_.T @ linalg_inv(precision) @ components_ precision /= -(self.noise_variance_**2) - _fill_or_add_to_diagonal(precision, 1.0 / self.noise_variance_, xp) + precision = _fill_or_add_to_diagonal(precision, 1.0 / self.noise_variance_, xp) return precision @abstractmethod diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py index f0e6cee65bc28..14e84d4f7331e 100644 --- a/sklearn/metrics/pairwise.py +++ b/sklearn/metrics/pairwise.py @@ -433,7 +433,7 @@ def _euclidean_distances(X, Y, X_norm_squared=None, Y_norm_squared=None, squared # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. if X is Y: - _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False) + distances = _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False) if squared: return distances @@ -1171,7 +1171,7 @@ def cosine_distances(X, Y=None): if X is Y or Y is None: # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. - _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False) + S = _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False) return S diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 1603ee1fc4211..fe9fbd351fe7b 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -557,7 +557,9 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): array_flat[:end:step] += value else: array_flat[:end:step] = value - print(f"{array_flat=}\n{array=}") + # When `array` is non-contiguous (e.g., after `transpose`), `reshape` must create + # a copy, and cannot return a view. Thus we need to *return* reshaped `array_flat`. + return xp.reshape(array_flat, array.shape) def _is_xp_namespace(xp, name): diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index 38a87418c5357..4002ec7c27e6f 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -583,39 +583,22 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): xp = _array_api_for_tests(array_namespace, device_) array_np = numpy.zeros((5, 4), dtype=numpy.int64) array_xp = xp.asarray(array_np) - _fill_or_add_to_diagonal(array_xp, value=1, xp=xp, add_value=False, wrap=wrap) + array_xp = _fill_or_add_to_diagonal( + array_xp, value=1, xp=xp, add_value=False, wrap=wrap + ) numpy.fill_diagonal(array_np, val=1, wrap=wrap) assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) -from sklearn.utils import gen_even_slices -from sklearn.utils.parallel import Parallel, delayed -from sklearn.metrics import euclidean_distances -from sklearn.utils._array_api import _find_matching_floating_dtype - -@pytest.mark.parametrize( - "array_namespace, device, dtype_name", - yield_namespace_device_dtype_combinations(), - ids=_get_namespace_device_dtype_ids, -) -def test_fill_or_add_to_diagonal_parallel(array_namespace, device, dtype_name): - """""" - xp = _array_api_for_tests(array_namespace, device) - n_samples = 10 - rng = numpy.random.RandomState(0) - X_np = numpy.array(5 * rng.random_sample((n_samples, n_samples)), dtype=dtype_name) - X_xp = xp.asarray(X_np, device=device) - ret_xp = xp.ones((n_samples, n_samples), dtype=X_xp.dtype, device=device).T - - def dumm_func(return_array, slice_): - return_array[..., slice_] = 1 - - fd = delayed(dumm_func) - with config_context(array_api_dispatch=True): - Parallel(backend="threading", n_jobs=2)( - fd(ret_xp, s) for s in gen_even_slices(n_samples, 2) - ) - _fill_or_add_to_diagonal(ret_xp, value=99, xp=xp, add_value=False) +def test_fill_or_add_to_diagonal_transpose(): + """Check `_fill_or_add_to_diagonal` when `reshape` returns a copy.""" + xp = _array_api_for_tests("numpy", None) + # Transposing an array makes it non-contiguous, meaning `reshape`, used within + # `_fill_or_add_to_diagonal`, returns a copy instead of a view. Note + # `numpy.fill_diagonal` avoids this problem as it uses `.flat` instead of `reshape` + array = numpy.ones((2, 2)).T + array = _fill_or_add_to_diagonal(array, value=99, xp=xp, add_value=False) + assert array[0, 0] == 99 @pytest.mark.parametrize("csr_container", CSR_CONTAINERS) From 2f2fcec29f74fb5635d8a2eeeccc84576e3a8a27 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 29 May 2025 12:42:30 +1000 Subject: [PATCH 05/22] Update sklearn/utils/tests/test_array_api.py Co-authored-by: Thomas J. Fan --- sklearn/utils/tests/test_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index d79937e0d9cb2..babb5ac0d74c3 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -597,7 +597,7 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): def test_fill_or_add_to_diagonal_transpose(): """Check `_fill_or_add_to_diagonal` when `reshape` returns a copy.""" xp = _array_api_for_tests("numpy", None) - # Transposing an array makes it non-contiguous, meaning `reshape`, used within + # Transposing an array makes it F-contiguous, meaning `reshape(x, (-1,))`, used within # `_fill_or_add_to_diagonal`, returns a copy instead of a view. Note # `numpy.fill_diagonal` avoids this problem as it uses `.flat` instead of `reshape` array = numpy.ones((2, 2)).T From e44e5277fcd337d1d79499e410cf20282890985c Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 29 May 2025 12:42:37 +1000 Subject: [PATCH 06/22] Update sklearn/utils/_array_api.py Co-authored-by: Thomas J. Fan --- sklearn/utils/_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index fe9fbd351fe7b..07eab7e1a2f2b 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -557,7 +557,7 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): array_flat[:end:step] += value else: array_flat[:end:step] = value - # When `array` is non-contiguous (e.g., after `transpose`), `reshape` must create + # When `array` is not C-contiguous, `reshape` must create # a copy, and cannot return a view. Thus we need to *return* reshaped `array_flat`. return xp.reshape(array_flat, array.shape) From e7ceb5334226c60eb0e012c773084f739025964f Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 29 May 2025 12:43:34 +1000 Subject: [PATCH 07/22] format --- sklearn/utils/tests/test_array_api.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index babb5ac0d74c3..b1647c4b536ce 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -597,8 +597,8 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): def test_fill_or_add_to_diagonal_transpose(): """Check `_fill_or_add_to_diagonal` when `reshape` returns a copy.""" xp = _array_api_for_tests("numpy", None) - # Transposing an array makes it F-contiguous, meaning `reshape(x, (-1,))`, used within - # `_fill_or_add_to_diagonal`, returns a copy instead of a view. Note + # Transposing an array makes it F-contiguous, meaning `reshape(x, (-1,))`, used + # within `_fill_or_add_to_diagonal`, returns a copy instead of a view. Note # `numpy.fill_diagonal` avoids this problem as it uses `.flat` instead of `reshape` array = numpy.ones((2, 2)).T array = _fill_or_add_to_diagonal(array, value=99, xp=xp, add_value=False) From 7e59fe6f032a8f3ff221ec16fee5cbaa84150c12 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 29 May 2025 18:49:12 +1000 Subject: [PATCH 08/22] review --- sklearn/utils/_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 07eab7e1a2f2b..88c702d172aa5 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -557,7 +557,7 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): array_flat[:end:step] += value else: array_flat[:end:step] = value - # When `array` is not C-contiguous, `reshape` must create + # When `array` is not C-contiguous, `reshape` creates # a copy, and cannot return a view. Thus we need to *return* reshaped `array_flat`. return xp.reshape(array_flat, array.shape) From 86bde6ca7ad759ac9d6e2dd8e2152e52a9a2f767 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 29 May 2025 18:49:39 +1000 Subject: [PATCH 09/22] format --- sklearn/utils/_array_api.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 88c702d172aa5..466880e2294ae 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -557,8 +557,8 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): array_flat[:end:step] += value else: array_flat[:end:step] = value - # When `array` is not C-contiguous, `reshape` creates - # a copy, and cannot return a view. Thus we need to *return* reshaped `array_flat`. + # When `array` is not C-contiguous, `reshape` creates a copy, and cannot + # return a view. Thus we need to *return* reshaped `array_flat`. return xp.reshape(array_flat, array.shape) From bc75c428cb68e74aa205696ca42556661a0c11e3 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Wed, 4 Jun 2025 21:05:11 +1000 Subject: [PATCH 10/22] wip --- sklearn/utils/_array_api.py | 86 ++++++++++++++++++++++++++- sklearn/utils/tests/test_array_api.py | 12 +++- 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 466880e2294ae..5a13617b2e76c 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -527,7 +527,7 @@ def _expit(X, xp=None): return 1.0 / (1.0 + xp.exp(-X)) -def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): +def _fill_or_add_to_diagonal_reshape(array, value, xp, add_value=True, wrap=False): """Implementation to facilitate adding or assigning specified values to the diagonal of a 2-d array. @@ -562,6 +562,90 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): return xp.reshape(array_flat, array.shape) +def _fill_diagonal_helper(array, value, xp, assignment_function): + """Implementation to facilitate adding or assigning specified values to the + diagonal of a 2-d array. + + The implementation is inspired from the `numpy.fill_diagonal` function: + https://github.com/numpy/numpy/blob/v2.0.0/numpy/lib/_index_tricks_impl.py#L799-L929 + """ + if array.ndim != 2: + raise ValueError( + f"array should be 2-d. Got array with shape {tuple(array.shape)}" + ) + + value = xp.asarray(value, dtype=array.dtype, device=device(array)) + if value.ndim not in [0, 1]: + raise ValueError( + "value needs to be a scalar or a 1d array, " + f"got a {value.ndim}d array instead" + ) + + step = array.shape[1] + 1 + # 'end' avoids wrapping in case of array is non-square with n_rows > n_columns + end = array.shape[1] * array.shape[1] + min_rows_columns = min(array.shape) + + if _is_numpy_namespace(xp): + assignment_function(array.flat, slice(None, end, step), value) + elif value.ndim == 1: + for i in range(min_rows_columns): + assignment_function(array, (i, i), value[i]) + else: + for i in range(min_rows_columns): + assignment_function(array, (i, i), value) + + +def _fill_diagonal_using_helper(array, value, xp): + def assignment_function(lhs, index, rhs): + lhs[index] = rhs + + return _fill_diagonal_helper(array, value, xp, assignment_function) + + +def _add_to_diagonal_using_helper(array, value, xp): + def assignment_function(lhs, index, rhs): + lhs[index] += rhs + + return _fill_diagonal_helper(array, value, xp, assignment_function) + + +def _fill_or_add_to_diagonal_loop(array, value, xp, add_value=True, wrap=False): + if _is_numpy_namespace(xp): + if add_value: + numpy.fill_diagonal(array, numpy.diagonal(array) + value, wrap=wrap) + else: + numpy.fill_diagonal(array, value, wrap=wrap) + return + + if wrap: + # Note `diagonal` does not support wrap, so difficult to implement + # `add_diagonal` with wrap + raise ValueError("`wrap=True` is not supported") + + if array.ndim != 2: + raise ValueError( + f"array should be 2-d. Got array with shape {tuple(array.shape)}" + ) + + value = xp.asarray(value, dtype=array.dtype, device=device(array)) + if value.ndim not in [0, 1]: + raise ValueError( + "value needs to be a scalar or a 1d array, " + f"got a {value.ndim}d array instead" + ) + if add_value: + value = xp.diagonal(array) + value + else: + if value.ndim == 0: + value = xp.repeat(value, min(array.shape)) + + min_rows_columns = min(array.shape) + + for i in range(min_rows_columns): + array[i, i] = value[i] + + def _is_xp_namespace(xp, name): return xp.__name__ in ( name, diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index b1647c4b536ce..d625e4b27394a 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -8,13 +8,15 @@ from sklearn._config import config_context from sklearn.base import BaseEstimator from sklearn.utils._array_api import ( + _add_to_diagonal_using_helper, _asarray_with_order, _atol_for_type, _average, _convert_to_numpy, _count_nonzero, _estimator_with_converted_arrays, - _fill_or_add_to_diagonal, + _fill_diagonal_using_helper, + _fill_or_add_to_diagonal_loop, _get_namespace_device_dtype_ids, _is_numpy_namespace, _isin, @@ -594,6 +596,14 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) +def test_aa(): + xp = _array_api_for_tests("numpy", None) + array_np = numpy.zeros((5, 4)) + _fill_or_add_to_diagonal_loop(array_np, 1, xp) + _fill_diagonal_using_helper(array_np, 1, xp) + _add_to_diagonal_using_helper(array_np, 1, xp) + + def test_fill_or_add_to_diagonal_transpose(): """Check `_fill_or_add_to_diagonal` when `reshape` returns a copy.""" xp = _array_api_for_tests("numpy", None) From ad49cc4ca844fffb8bbcdeea4444e05089b39e2a Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 5 Jun 2025 21:27:25 +1000 Subject: [PATCH 11/22] wip --- sklearn/decomposition/_base.py | 6 +- sklearn/metrics/pairwise.py | 4 +- sklearn/utils/_array_api.py | 112 ++++++++++++-------------- sklearn/utils/tests/test_array_api.py | 18 +---- 4 files changed, 61 insertions(+), 79 deletions(-) diff --git a/sklearn/decomposition/_base.py b/sklearn/decomposition/_base.py index edbfd1d019cd4..783c316b50f27 100644 --- a/sklearn/decomposition/_base.py +++ b/sklearn/decomposition/_base.py @@ -47,7 +47,7 @@ def get_covariance(self): xp.asarray(0.0, device=device(exp_var), dtype=exp_var.dtype), ) cov = (components_.T * exp_var_diff) @ components_ - cov = _fill_or_add_to_diagonal(cov, self.noise_variance_, xp) + _fill_or_add_to_diagonal(cov, self.noise_variance_, xp) return cov def get_precision(self): @@ -89,10 +89,10 @@ def get_precision(self): xp.asarray(0.0, device=device(exp_var)), ) precision = components_ @ components_.T / self.noise_variance_ - precision = _fill_or_add_to_diagonal(precision, 1.0 / exp_var_diff, xp) + _fill_or_add_to_diagonal(precision, 1.0 / exp_var_diff, xp) precision = components_.T @ linalg_inv(precision) @ components_ precision /= -(self.noise_variance_**2) - precision = _fill_or_add_to_diagonal(precision, 1.0 / self.noise_variance_, xp) + _fill_or_add_to_diagonal(precision, 1.0 / self.noise_variance_, xp) return precision @abstractmethod diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py index 14e84d4f7331e..f0e6cee65bc28 100644 --- a/sklearn/metrics/pairwise.py +++ b/sklearn/metrics/pairwise.py @@ -433,7 +433,7 @@ def _euclidean_distances(X, Y, X_norm_squared=None, Y_norm_squared=None, squared # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. if X is Y: - distances = _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False) + _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False) if squared: return distances @@ -1171,7 +1171,7 @@ def cosine_distances(X, Y=None): if X is Y or Y is None: # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. - S = _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False) + _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False) return S diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 5a13617b2e76c..45f30575eb75b 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -527,39 +527,69 @@ def _expit(X, xp=None): return 1.0 / (1.0 + xp.exp(-X)) -def _fill_or_add_to_diagonal_reshape(array, value, xp, add_value=True, wrap=False): - """Implementation to facilitate adding or assigning specified values to the - diagonal of a 2-d array. - - If ``add_value`` is `True` then the values will be added to the diagonal - elements otherwise the values will be assigned to the diagonal elements. - By default, ``add_value`` is set to `True. This is currently only - supported for 2-d arrays. +def _fill_or_add_to_diagonal(array, value, xp, add_value=False): + """Minimal implementation of `numpy.fill_diagonal`, which - The implementation is taken from the `numpy.fill_diagonal` function: - https://github.com/numpy/numpy/blob/v2.0.0/numpy/lib/_index_tricks_impl.py#L799-L929 + `wrap` is not supported (i.e. always False). `value` should be a scalar or + 1D of greater or equal length as the diagonal (i.e., `value` is never repeated + when shorter). """ if array.ndim != 2: raise ValueError( - f"array should be 2-d. Got array with shape {tuple(array.shape)}" + f"`array` should be 2D. Got array with shape {tuple(array.shape)}" ) value = xp.asarray(value, dtype=array.dtype, device=device(array)) - end = None - # Explicit, fast formula for the common case. For 2-d arrays, we - # accept rectangular ones. - step = array.shape[1] + 1 - if not wrap: - end = array.shape[1] * array.shape[1] + if value.ndim not in [0, 1]: + raise ValueError( + "value needs to be a scalar or a 1d array, " + f"got a {value.ndim}d array instead" + ) - array_flat = xp.reshape(array, (-1,)) + min_rows_columns = min(array.shape) if add_value: - array_flat[:end:step] += value + value = xp.diagonal(array) + value + if value.ndim == 0: + for i in range(min_rows_columns): + array[i, i] = value else: - array_flat[:end:step] = value - # When `array` is not C-contiguous, `reshape` creates a copy, and cannot - # return a view. Thus we need to *return* reshaped `array_flat`. - return xp.reshape(array_flat, array.shape) + for i in range(min_rows_columns): + array[i, i] = value[i] + + +# def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): +# """Implementation to facilitate adding or assigning specified values to the +# diagonal of a 2-d array. + +# If ``add_value`` is `True` then the values will be added to the diagonal +# elements otherwise the values will be assigned to the diagonal elements. +# By default, ``add_value`` is set to `True. This is currently only +# supported for 2-d arrays. + +# The implementation is taken from the `numpy.fill_diagonal` function: +# https://github.com/numpy/numpy/blob/v2.0.0/numpy/lib/_index_tricks_impl.py#L799-L929 +# """ +# if array.ndim != 2: +# raise ValueError( +# f"array should be 2-d. Got array with shape {tuple(array.shape)}" +# ) + +# value = xp.asarray(value, dtype=array.dtype, device=device(array)) +# end = None +# # Explicit, fast formula for the common case. For 2-d arrays, we +# # accept rectangular ones. +# step = array.shape[1] + 1 +# if not wrap: +# end = array.shape[1] * array.shape[1] + +# array_flat = xp.reshape(array, (-1,)) +# if add_value: +# array_flat[:end:step] += value +# else: +# array_flat[:end:step] = value +# # When `array` is not C-contiguous, `reshape` creates a copy, and cannot +# # return a view. Thus we need to *return* reshaped `array_flat`. +# return xp.reshape(array_flat, array.shape) def _fill_diagonal_helper(array, value, xp, assignment_function): @@ -610,42 +640,6 @@ def assignment_function(lhs, index, rhs): return _fill_diagonal_helper(array, value, xp, assignment_function) -def _fill_or_add_to_diagonal_loop(array, value, xp, add_value=True, wrap=False): - if _is_numpy_namespace(xp): - if add_value: - numpy.fill_diagonal(array, numpy.diagonal(array) + value, wrap=wrap) - else: - numpy.fill_diagonal(array, value, wrap=wrap) - return - - if wrap: - # Note `diagonal` does not support wrap, so difficult to implement - # `add_diagonal` with wrap - raise ValueError("`wrap=True` is not supported") - - if array.ndim != 2: - raise ValueError( - f"array should be 2-d. Got array with shape {tuple(array.shape)}" - ) - - value = xp.asarray(value, dtype=array.dtype, device=device(array)) - if value.ndim not in [0, 1]: - raise ValueError( - "value needs to be a scalar or a 1d array, " - f"got a {value.ndim}d array instead" - ) - if add_value: - value = xp.diagonal(array) + value - else: - if value.ndim == 0: - value = xp.repeat(value, min(array.shape)) - - min_rows_columns = min(array.shape) - - for i in range(min_rows_columns): - array[i, i] = value[i] - - def _is_xp_namespace(xp, name): return xp.__name__ in ( name, diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index d625e4b27394a..23a7921063fd7 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -8,15 +8,13 @@ from sklearn._config import config_context from sklearn.base import BaseEstimator from sklearn.utils._array_api import ( - _add_to_diagonal_using_helper, _asarray_with_order, _atol_for_type, _average, _convert_to_numpy, _count_nonzero, _estimator_with_converted_arrays, - _fill_diagonal_using_helper, - _fill_or_add_to_diagonal_loop, + _fill_or_add_to_diagonal, _get_namespace_device_dtype_ids, _is_numpy_namespace, _isin, @@ -589,21 +587,11 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): numpy.fill_diagonal(array_np, val=1, wrap=wrap) with config_context(array_api_dispatch=True): - array_xp = _fill_or_add_to_diagonal( - array_xp, value=1, xp=xp, add_value=False, wrap=wrap - ) + _fill_or_add_to_diagonal(array_xp, value=1, xp=xp, add_value=False, wrap=wrap) assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) -def test_aa(): - xp = _array_api_for_tests("numpy", None) - array_np = numpy.zeros((5, 4)) - _fill_or_add_to_diagonal_loop(array_np, 1, xp) - _fill_diagonal_using_helper(array_np, 1, xp) - _add_to_diagonal_using_helper(array_np, 1, xp) - - def test_fill_or_add_to_diagonal_transpose(): """Check `_fill_or_add_to_diagonal` when `reshape` returns a copy.""" xp = _array_api_for_tests("numpy", None) @@ -611,7 +599,7 @@ def test_fill_or_add_to_diagonal_transpose(): # within `_fill_or_add_to_diagonal`, returns a copy instead of a view. Note # `numpy.fill_diagonal` avoids this problem as it uses `.flat` instead of `reshape` array = numpy.ones((2, 2)).T - array = _fill_or_add_to_diagonal(array, value=99, xp=xp, add_value=False) + _fill_or_add_to_diagonal(array, value=99, xp=xp, add_value=False) assert array[0, 0] == 99 From 2e10b83eabab8a618ef1866780b7a7ba62bd449b Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Thu, 5 Jun 2025 21:46:56 +1000 Subject: [PATCH 12/22] add back wrap to make tests pass --- sklearn/utils/_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 45f30575eb75b..2769bc5c00979 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -527,7 +527,7 @@ def _expit(X, xp=None): return 1.0 / (1.0 + xp.exp(-X)) -def _fill_or_add_to_diagonal(array, value, xp, add_value=False): +def _fill_or_add_to_diagonal(array, value, xp, add_value=False, wrap=False): """Minimal implementation of `numpy.fill_diagonal`, which `wrap` is not supported (i.e. always False). `value` should be a scalar or From 0d0b4cbd6793af2b6e7c2243ad63b57f97b2243a Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Fri, 6 Jun 2025 21:17:43 +1000 Subject: [PATCH 13/22] fix default add_value True --- sklearn/utils/_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 2769bc5c00979..ca88368d26058 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -527,7 +527,7 @@ def _expit(X, xp=None): return 1.0 / (1.0 + xp.exp(-X)) -def _fill_or_add_to_diagonal(array, value, xp, add_value=False, wrap=False): +def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): """Minimal implementation of `numpy.fill_diagonal`, which `wrap` is not supported (i.e. always False). `value` should be a scalar or From 7c237dc440ff659ca21231dd05a9ea60a1278e1e Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Sat, 7 Jun 2025 21:55:18 +1000 Subject: [PATCH 14/22] linalg diagonal --- sklearn/utils/_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index ca88368d26058..c5a7d303459ef 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -548,7 +548,7 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): min_rows_columns = min(array.shape) if add_value: - value = xp.diagonal(array) + value + value = xp.linalg.diagonal(array) + value if value.ndim == 0: for i in range(min_rows_columns): array[i, i] = value From 47824d88f5342e5c352d58527dabf9e401d1d22f Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Fri, 13 Jun 2025 13:24:49 +1000 Subject: [PATCH 15/22] wip --- sklearn/utils/_array_api.py | 126 ++++++++------------------ sklearn/utils/tests/test_array_api.py | 56 +++++++++--- 2 files changed, 80 insertions(+), 102 deletions(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index c5a7d303459ef..788d808e74057 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -527,13 +527,8 @@ def _expit(X, xp=None): return 1.0 / (1.0 + xp.exp(-X)) -def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): - """Minimal implementation of `numpy.fill_diagonal`, which - - `wrap` is not supported (i.e. always False). `value` should be a scalar or - 1D of greater or equal length as the diagonal (i.e., `value` is never repeated - when shorter). - """ +def _validate_diagonal_args(array, value, xp): + """Helper to validate arguments to `_fill_diagonal`/`add_to_diagonal`.""" if array.ndim != 2: raise ValueError( f"`array` should be 2D. Got array with shape {tuple(array.shape)}" @@ -542,102 +537,55 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): value = xp.asarray(value, dtype=array.dtype, device=device(array)) if value.ndim not in [0, 1]: raise ValueError( - "value needs to be a scalar or a 1d array, " - f"got a {value.ndim}d array instead" + "`value` needs to be a scalar or a 1D array, " + f"got a {value.ndim}D array instead." ) - min_rows_columns = min(array.shape) - if add_value: - value = xp.linalg.diagonal(array) + value - if value.ndim == 0: - for i in range(min_rows_columns): - array[i, i] = value - else: - for i in range(min_rows_columns): - array[i, i] = value[i] - - -# def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False): -# """Implementation to facilitate adding or assigning specified values to the -# diagonal of a 2-d array. - -# If ``add_value`` is `True` then the values will be added to the diagonal -# elements otherwise the values will be assigned to the diagonal elements. -# By default, ``add_value`` is set to `True. This is currently only -# supported for 2-d arrays. - -# The implementation is taken from the `numpy.fill_diagonal` function: -# https://github.com/numpy/numpy/blob/v2.0.0/numpy/lib/_index_tricks_impl.py#L799-L929 -# """ -# if array.ndim != 2: -# raise ValueError( -# f"array should be 2-d. Got array with shape {tuple(array.shape)}" -# ) - -# value = xp.asarray(value, dtype=array.dtype, device=device(array)) -# end = None -# # Explicit, fast formula for the common case. For 2-d arrays, we -# # accept rectangular ones. -# step = array.shape[1] + 1 -# if not wrap: -# end = array.shape[1] * array.shape[1] - -# array_flat = xp.reshape(array, (-1,)) -# if add_value: -# array_flat[:end:step] += value -# else: -# array_flat[:end:step] = value -# # When `array` is not C-contiguous, `reshape` creates a copy, and cannot -# # return a view. Thus we need to *return* reshaped `array_flat`. -# return xp.reshape(array_flat, array.shape) - - -def _fill_diagonal_helper(array, value, xp, assignment_function): - """Implementation to facilitate adding or assigning specified values to the - diagonal of a 2-d array. - - The implementation is inspired from the `numpy.fill_diagonal` function: - https://github.com/numpy/numpy/blob/v2.0.0/numpy/lib/_index_tricks_impl.py#L799-L929 - """ - if array.ndim != 2: + if value.ndim == 1 and value.shape[0] != min_rows_columns: raise ValueError( - f"array should be 2-d. Got array with shape {tuple(array.shape)}" + "`value` needs to be a scalar or 1D array of the same length as the " + f"diagonal of `array` ({min_rows_columns}). Got {value.shape[0]}" ) - value = xp.asarray(value, dtype=array.dtype, device=device(array)) - if value.ndim not in [0, 1]: - raise ValueError( - "value needs to be a scalar or a 1d array, " - f"got a {value.ndim}d array instead" - ) + return value, min_rows_columns - step = array.shape[1] + 1 - # 'end' avoids wrapping in case of array is non-square with n_rows > n_columns - end = array.shape[1] * array.shape[1] - min_rows_columns = min(array.shape) + +def _fill_diagonal(array, value, xp): + """Minimal implementation of `numpy.fill_diagonal`. + + `wrap` is not supported (i.e. always False). `value` should be a scalar or + 1D of greater or equal length as the diagonal (i.e., `value` is never repeated + when shorter). + + Note `array` is altered in place. + """ + value, min_rows_columns = _validate_diagonal_args(array, value, xp) if _is_numpy_namespace(xp): - assignment_function(array.flat, slice(None, end, step), value) - elif value.ndim == 1: - for i in range(min_rows_columns): - assignment_function(array, (i, i), value[i]) + xp.fill_diagonal(array, value, wrap=False) else: - for i in range(min_rows_columns): - assignment_function(array, (i, i), value) - + if value.ndim == 0: + for i in range(min_rows_columns): + array[i, i] = value + else: + for i in range(min_rows_columns): + array[i, i] = value[i] -def _fill_diagonal_using_helper(array, value, xp): - def assignment_function(lhs, index, rhs): - lhs[index] = rhs - return _fill_diagonal_helper(array, value, xp, assignment_function) +def _add_to_diagonal(array, value, xp): + """Add `value` to diagonal of `array`. + Related to `fill_diagonal`. `value` should be a scalar or + 1D of greater or equal length as the diagonal (i.e., `value` is never repeated + when shorter). -def _add_to_diagonal_using_helper(array, value, xp): - def assignment_function(lhs, index, rhs): - lhs[index] += rhs + Note `array` is altered in place. + """ + value, min_rows_columns = _validate_diagonal_args(array, value, xp) - return _fill_diagonal_helper(array, value, xp, assignment_function) + value = xp.linalg.diagonal(array) + value + for i in range(min_rows_columns): + array[i, i] = value[i] def _is_xp_namespace(xp, name): diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index 23a7921063fd7..eac8ba67c7a6d 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -8,13 +8,14 @@ from sklearn._config import config_context from sklearn.base import BaseEstimator from sklearn.utils._array_api import ( + _add_to_diagonal, _asarray_with_order, _atol_for_type, _average, _convert_to_numpy, _count_nonzero, _estimator_with_converted_arrays, - _fill_or_add_to_diagonal, + _fill_diagonal, _get_namespace_device_dtype_ids, _is_numpy_namespace, _isin, @@ -573,12 +574,52 @@ def test_count_nonzero( assert device(array_xp) == device(result) +def test_validate_diagonal_args(): + """Check `_validate_diagonal_args` raises the correct errors.""" + + + + +@pytest.mark.parametrize("array, c_contiguity", [(numpy.zeros((3,4)), True), (numpy.zeros((3,4)).T, False)]) +def test_fill_diagonal(array, c_contiguity): + """Check `_fill_diagonal` behaviour is correct with numpy arrays.""" + xp = _array_api_for_tests("numpy", None) + assert array.flags['C_CONTIGUOUS'] == c_contiguity + + _fill_diagonal(array, 1, xp) + assert_allclose(array.diagonal(), numpy.ones((3,))) + + _fill_diagonal(array, [0, 1, 2], xp) + assert_allclose(array.diagonal(), numpy.arange(3)) + + fill_array = numpy.array([11, 12, 13]) + _fill_diagonal(array, fill_array, xp) + assert_allclose(array.diagonal(), fill_array) + + +@pytest.mark.parametrize("array, c_contiguity", [(numpy.zeros((3,4)), True), (numpy.zeros((3,4)).T, False)]) +def test_add_to_diagonal(array, c_contiguity): + """Check `_add_to_diagonal` behaviour is correct with numpy arrays.""" + xp = _array_api_for_tests("numpy", None) + assert array.flags['C_CONTIGUOUS'] == c_contiguity + + _add_to_diagonal(array, 1, xp) + assert_allclose(array.diagonal(), numpy.ones((3,))) + + _add_to_diagonal(array, [0, 1, 2], xp) + assert_allclose(array.diagonal(), numpy.arange(3)) + + fill_array = numpy.array([11, 12, 13]) + _add_to_diagonal(array, fill_array, xp) + assert_allclose(array.diagonal(), fill_array) + + + @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) -@pytest.mark.parametrize("wrap", [True, False]) def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): xp = _array_api_for_tests(array_namespace, device_) @@ -592,17 +633,6 @@ def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) -def test_fill_or_add_to_diagonal_transpose(): - """Check `_fill_or_add_to_diagonal` when `reshape` returns a copy.""" - xp = _array_api_for_tests("numpy", None) - # Transposing an array makes it F-contiguous, meaning `reshape(x, (-1,))`, used - # within `_fill_or_add_to_diagonal`, returns a copy instead of a view. Note - # `numpy.fill_diagonal` avoids this problem as it uses `.flat` instead of `reshape` - array = numpy.ones((2, 2)).T - _fill_or_add_to_diagonal(array, value=99, xp=xp, add_value=False) - assert array[0, 0] == 99 - - @pytest.mark.parametrize("csr_container", CSR_CONTAINERS) @pytest.mark.parametrize("dispatch", [True, False]) def test_sparse_device(csr_container, dispatch): From 77869ded13541722d2bdd5345c672bd6a3616a7e Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Fri, 13 Jun 2025 13:37:09 +1000 Subject: [PATCH 16/22] wip test --- sklearn/utils/tests/test_array_api.py | 45 ++++++++++++++++++++------- 1 file changed, 33 insertions(+), 12 deletions(-) diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index eac8ba67c7a6d..8ad143755ff76 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -24,6 +24,7 @@ _nanmean, _nanmin, _ravel, + _validate_diagonal_args, device, get_namespace, get_namespace_and_device, @@ -574,26 +575,46 @@ def test_count_nonzero( assert device(array_xp) == device(result) -def test_validate_diagonal_args(): +@pytest.mark.parametrize( + "array, value, match", + [ + (numpy.array([1, 2, 3]), 1, "`array` should be 2D"), + (numpy.array([[1, 2], [3, 4]]), numpy.array([1, 2, 3]), "`value` needs to be"), + (numpy.array([[1, 2], [3, 4]]), [1, 2, 3], "`value` needs to be"), + (numpy.array([[1, 2], [3, 4]]), numpy.array([[1, 2], [3, 4]]), "`value` needs to be a"), + ], +) +def test_validate_diagonal_args(array, value, match): """Check `_validate_diagonal_args` raises the correct errors.""" + xp = _array_api_for_tests("numpy", None) + with pytest.raises(ValueError, match=match): + _validate_diagonal_args(array, value, xp) - - -@pytest.mark.parametrize("array, c_contiguity", [(numpy.zeros((3,4)), True), (numpy.zeros((3,4)).T, False)]) -def test_fill_diagonal(array, c_contiguity): - """Check `_fill_diagonal` behaviour is correct with numpy arrays.""" +@pytest.mark.parametrize("function", ["fill", "add"]) +@pytest.mark.parametrize("c_contiguity", [True, False]) +def test_fill_and_add_to_diagonal(c_contiguity, function): + """Check `_fill/add_to_diagonal` behaviour correct with numpy arrays.""" xp = _array_api_for_tests("numpy", None) + if c_contiguity: + array = numpy.zeros((3,4)) + else: + array = numpy.zeros((3,4)).T assert array.flags['C_CONTIGUOUS'] == c_contiguity - _fill_diagonal(array, 1, xp) + if function == "fill": + func = _fill_diagonal + else: + func = _add_to_diagonal + + func(array, 1, xp) assert_allclose(array.diagonal(), numpy.ones((3,))) - _fill_diagonal(array, [0, 1, 2], xp) + func(array, [0, 1, 2], xp) assert_allclose(array.diagonal(), numpy.arange(3)) fill_array = numpy.array([11, 12, 13]) - _fill_diagonal(array, fill_array, xp) + func(array, fill_array, xp) assert_allclose(array.diagonal(), fill_array) @@ -609,9 +630,9 @@ def test_add_to_diagonal(array, c_contiguity): _add_to_diagonal(array, [0, 1, 2], xp) assert_allclose(array.diagonal(), numpy.arange(3)) - fill_array = numpy.array([11, 12, 13]) - _add_to_diagonal(array, fill_array, xp) - assert_allclose(array.diagonal(), fill_array) + add_array = numpy.array([11, 12, 13]) + _add_to_diagonal(array, add_array, xp) + assert_allclose(array.diagonal(), add_array) From 958e18a12531d8930ef0d25515eefed128367dd1 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Fri, 13 Jun 2025 13:53:25 +1000 Subject: [PATCH 17/22] add test --- sklearn/decomposition/_base.py | 8 ++-- sklearn/metrics/pairwise.py | 6 +-- sklearn/utils/tests/test_array_api.py | 62 +++++++++++++++++---------- 3 files changed, 46 insertions(+), 30 deletions(-) diff --git a/sklearn/decomposition/_base.py b/sklearn/decomposition/_base.py index 783c316b50f27..85cc746fd9b8a 100644 --- a/sklearn/decomposition/_base.py +++ b/sklearn/decomposition/_base.py @@ -9,7 +9,7 @@ from scipy import linalg from ..base import BaseEstimator, ClassNamePrefixFeaturesOutMixin, TransformerMixin -from ..utils._array_api import _fill_or_add_to_diagonal, device, get_namespace +from ..utils._array_api import _add_to_diagonal, device, get_namespace from ..utils.validation import check_is_fitted, validate_data @@ -47,7 +47,7 @@ def get_covariance(self): xp.asarray(0.0, device=device(exp_var), dtype=exp_var.dtype), ) cov = (components_.T * exp_var_diff) @ components_ - _fill_or_add_to_diagonal(cov, self.noise_variance_, xp) + _add_to_diagonal(cov, self.noise_variance_, xp) return cov def get_precision(self): @@ -89,10 +89,10 @@ def get_precision(self): xp.asarray(0.0, device=device(exp_var)), ) precision = components_ @ components_.T / self.noise_variance_ - _fill_or_add_to_diagonal(precision, 1.0 / exp_var_diff, xp) + _add_to_diagonal(precision, 1.0 / exp_var_diff, xp) precision = components_.T @ linalg_inv(precision) @ components_ precision /= -(self.noise_variance_**2) - _fill_or_add_to_diagonal(precision, 1.0 / self.noise_variance_, xp) + _add_to_diagonal(precision, 1.0 / self.noise_variance_, xp) return precision @abstractmethod diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py index f0e6cee65bc28..5dfcec2856bdb 100644 --- a/sklearn/metrics/pairwise.py +++ b/sklearn/metrics/pairwise.py @@ -19,7 +19,7 @@ from ..preprocessing import normalize from ..utils import check_array, gen_batches, gen_even_slices from ..utils._array_api import ( - _fill_or_add_to_diagonal, + _fill_diagonal, _find_matching_floating_dtype, _is_numpy_namespace, _max_precision_float_dtype, @@ -433,7 +433,7 @@ def _euclidean_distances(X, Y, X_norm_squared=None, Y_norm_squared=None, squared # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. if X is Y: - _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False) + _fill_diagonal(distances, 0, xp=xp, add_value=False) if squared: return distances @@ -1171,7 +1171,7 @@ def cosine_distances(X, Y=None): if X is Y or Y is None: # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. - _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False) + _fill_diagonal(S, 0.0, xp, add_value=False) return S diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index 8ad143755ff76..c43057ee947b9 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -581,7 +581,11 @@ def test_count_nonzero( (numpy.array([1, 2, 3]), 1, "`array` should be 2D"), (numpy.array([[1, 2], [3, 4]]), numpy.array([1, 2, 3]), "`value` needs to be"), (numpy.array([[1, 2], [3, 4]]), [1, 2, 3], "`value` needs to be"), - (numpy.array([[1, 2], [3, 4]]), numpy.array([[1, 2], [3, 4]]), "`value` needs to be a"), + ( + numpy.array([[1, 2], [3, 4]]), + numpy.array([[1, 2], [3, 4]]), + "`value` needs to be a", + ), ], ) def test_validate_diagonal_args(array, value, match): @@ -597,10 +601,10 @@ def test_fill_and_add_to_diagonal(c_contiguity, function): """Check `_fill/add_to_diagonal` behaviour correct with numpy arrays.""" xp = _array_api_for_tests("numpy", None) if c_contiguity: - array = numpy.zeros((3,4)) + array = numpy.zeros((3, 4)) else: - array = numpy.zeros((3,4)).T - assert array.flags['C_CONTIGUOUS'] == c_contiguity + array = numpy.zeros((3, 4)).T + assert array.flags["C_CONTIGUOUS"] == c_contiguity if function == "fill": func = _fill_diagonal @@ -611,29 +615,38 @@ def test_fill_and_add_to_diagonal(c_contiguity, function): assert_allclose(array.diagonal(), numpy.ones((3,))) func(array, [0, 1, 2], xp) - assert_allclose(array.diagonal(), numpy.arange(3)) + if function == "fill": + expected_diag = numpy.arange(3) + else: + expected_diag = numpy.ones((3,)) + numpy.arange(3) + assert_allclose(array.diagonal(), expected_diag) fill_array = numpy.array([11, 12, 13]) func(array, fill_array, xp) - assert_allclose(array.diagonal(), fill_array) - + if function == "fill": + expected_diag = fill_array + else: + expected_diag = fill_array + numpy.arange(3) + numpy.ones((3,)) + assert_allclose(array.diagonal(), expected_diag) -@pytest.mark.parametrize("array, c_contiguity", [(numpy.zeros((3,4)), True), (numpy.zeros((3,4)).T, False)]) -def test_add_to_diagonal(array, c_contiguity): - """Check `_add_to_diagonal` behaviour is correct with numpy arrays.""" - xp = _array_api_for_tests("numpy", None) - assert array.flags['C_CONTIGUOUS'] == c_contiguity - _add_to_diagonal(array, 1, xp) - assert_allclose(array.diagonal(), numpy.ones((3,))) +@pytest.mark.parametrize( + "array_namespace, device_, dtype_name", + yield_namespace_device_dtype_combinations(), + ids=_get_namespace_device_dtype_ids, +) +def test_fill_diagonal(array_namespace, device_, dtype_name): + """Check array API `_fill_diagonal` consistent with `numpy._fill_diagonal`.""" + xp = _array_api_for_tests(array_namespace, device_) - _add_to_diagonal(array, [0, 1, 2], xp) - assert_allclose(array.diagonal(), numpy.arange(3)) + array_np = numpy.zeros((3, 4), dtype=dtype_name) + array_xp = xp.asarray(array_np.copy(), device=device_) - add_array = numpy.array([11, 12, 13]) - _add_to_diagonal(array, add_array, xp) - assert_allclose(array.diagonal(), add_array) + numpy.fill_diagonal(array_np, val=1) + with config_context(array_api_dispatch=True): + _fill_diagonal(array_xp, value=1, xp=xp) + assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) @pytest.mark.parametrize( @@ -641,15 +654,18 @@ def test_add_to_diagonal(array, c_contiguity): yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) -def test_fill_or_add_to_diagonal(array_namespace, device_, dtype_name, wrap): +def test_add_to_diagonal(array_namespace, device_, dtype_name): + """Check `_add_to_diagonal` consistent between array API xp and numpy namespace.""" xp = _array_api_for_tests(array_namespace, device_) + np_xp = _array_api_for_tests("numpy", None) - array_np = numpy.zeros((5, 4), dtype=dtype_name) + array_np = numpy.zeros((3, 4), dtype=dtype_name) array_xp = xp.asarray(array_np.copy(), device=device_) - numpy.fill_diagonal(array_np, val=1, wrap=wrap) + add_val = [1, 2, 3] + _fill_diagonal(array_np, value=add_val, xp=np_xp) with config_context(array_api_dispatch=True): - _fill_or_add_to_diagonal(array_xp, value=1, xp=xp, add_value=False, wrap=wrap) + _fill_diagonal(array_xp, value=add_val, xp=xp) assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) From fe97489e2e643502f9f52c9075afd562d95af742 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Fri, 13 Jun 2025 14:07:47 +1000 Subject: [PATCH 18/22] fix param --- sklearn/metrics/pairwise.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py index 5dfcec2856bdb..d6ba6f53fa6af 100644 --- a/sklearn/metrics/pairwise.py +++ b/sklearn/metrics/pairwise.py @@ -433,7 +433,7 @@ def _euclidean_distances(X, Y, X_norm_squared=None, Y_norm_squared=None, squared # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. if X is Y: - _fill_diagonal(distances, 0, xp=xp, add_value=False) + _fill_diagonal(distances, 0, xp=xp) if squared: return distances @@ -1171,7 +1171,7 @@ def cosine_distances(X, Y=None): if X is Y or Y is None: # Ensure that distances between vectors and themselves are set to 0.0. # This may not be the case due to floating point rounding errors. - _fill_diagonal(S, 0.0, xp, add_value=False) + _fill_diagonal(S, 0.0, xp) return S From a36e0b16f3159b79919000db21dcb24924d86ec8 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Fri, 13 Jun 2025 15:55:20 +1000 Subject: [PATCH 19/22] use flat for add to --- sklearn/utils/_array_api.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 788d808e74057..35dc4fe65af6b 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -583,6 +583,13 @@ def _add_to_diagonal(array, value, xp): """ value, min_rows_columns = _validate_diagonal_args(array, value, xp) + if _is_numpy_namespace(xp): + step = array.shape[1] + 1 + # Ensure we do not wrap + end = array.shape[1] * array.shape[1] + array.flat[:end:step] += value + return + value = xp.linalg.diagonal(array) + value for i in range(min_rows_columns): array[i, i] = value[i] From a3d452d5ab50c4e16575bca25760ddbbde99f108 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Mon, 16 Jun 2025 11:59:34 +1000 Subject: [PATCH 20/22] improve test --- sklearn/utils/tests/test_array_api.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py index c43057ee947b9..4b9420bd366b0 100644 --- a/sklearn/utils/tests/test_array_api.py +++ b/sklearn/utils/tests/test_array_api.py @@ -630,17 +630,25 @@ def test_fill_and_add_to_diagonal(c_contiguity, function): assert_allclose(array.diagonal(), expected_diag) +@pytest.mark.parametrize("array", ["standard", "transposed", "non-contiguous"]) @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) -def test_fill_diagonal(array_namespace, device_, dtype_name): +def test_fill_diagonal(array, array_namespace, device_, dtype_name): """Check array API `_fill_diagonal` consistent with `numpy._fill_diagonal`.""" xp = _array_api_for_tests(array_namespace, device_) - - array_np = numpy.zeros((3, 4), dtype=dtype_name) - array_xp = xp.asarray(array_np.copy(), device=device_) + array_np = numpy.zeros((4, 5), dtype=dtype_name) + + if array == "transposed": + array_xp = xp.asarray(array_np.copy(), device=device_).T + array_np = array_np.T + elif array == "non-contiguous": + array_xp = xp.asarray(array_np.copy(), device=device_)[::2, ::2] + array_np = array_np[::2, ::2] + else: + array_xp = xp.asarray(array_np.copy(), device=device_) numpy.fill_diagonal(array_np, val=1) with config_context(array_api_dispatch=True): From 493b0f67a158cc1160b6eb57fb518f84a1663e8c Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Mon, 16 Jun 2025 12:09:42 +1000 Subject: [PATCH 21/22] add comment --- sklearn/utils/_array_api.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index 35dc4fe65af6b..b9e757782a2f9 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -564,6 +564,10 @@ def _fill_diagonal(array, value, xp): if _is_numpy_namespace(xp): xp.fill_diagonal(array, value, wrap=False) else: + # TODO: when array libraries support `reshape(copy)`, use + # `reshape(array, (-1,), copy=False)`, then fill with `[:end:step]` (within + # `try/except`). This is faster than for loop, when no copy needs to be + # made within `reshape`. See #31445 for details. if value.ndim == 0: for i in range(min_rows_columns): array[i, i] = value @@ -590,6 +594,10 @@ def _add_to_diagonal(array, value, xp): array.flat[:end:step] += value return + # TODO: when array libraries support `reshape(copy)`, use + # `reshape(array, (-1,), copy=False)`, then fill with `[:end:step]` (within + # `try/except`). This is faster than for loop, when no copy needs to be + # made within `reshape`. See #31445 for details. value = xp.linalg.diagonal(array) + value for i in range(min_rows_columns): array[i, i] = value[i] From a587d0d8fee31c814439925cacc7fbc1ebb8c8d4 Mon Sep 17 00:00:00 2001 From: Lucy Liu Date: Wed, 18 Jun 2025 11:06:06 +1000 Subject: [PATCH 22/22] review --- sklearn/utils/_array_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py index b9e757782a2f9..33e4f719007c3 100644 --- a/sklearn/utils/_array_api.py +++ b/sklearn/utils/_array_api.py @@ -528,7 +528,7 @@ def _expit(X, xp=None): def _validate_diagonal_args(array, value, xp): - """Helper to validate arguments to `_fill_diagonal`/`add_to_diagonal`.""" + """Validate arguments to `_fill_diagonal`/`_add_to_diagonal`.""" if array.ndim != 2: raise ValueError( f"`array` should be 2D. Got array with shape {tuple(array.shape)}"