Fix _fill_or_add_to_diagonal when reshape returns copy (scikit-learn#31445)

lucyleeow · thomasjpfan · web-flow · commit 0fc081a4e131 · 2025-06-19T16:48:50.000+02:00
Co-authored-by: Thomas J. Fan &lt;thomasjpfan@gmail.com&gt;
diff --git 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
@@ -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,
@@ -439,7 +439,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_diagonal(distances, 0, xp=xp)
 
     if squared:
         return distances
@@ -1177,7 +1177,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_diagonal(S, 0.0, xp)
     return S
 
 
@@ -1982,7 +1982,7 @@ def _parallel_pairwise(X, Y, func, n_jobs, **kwds):
     if (X is Y or Y is None) and func is euclidean_distances:
         # zeroing diagonal for euclidean norm.
         # TODO: do it also for other norms.
-        ret = _fill_or_add_to_diagonal(ret, 0, xp=xp, add_value=False)
+        _fill_diagonal(ret, 0, xp=xp)
 
     # Transform output back
     return ret.T
diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
@@ -527,40 +527,80 @@ 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):
-    """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
-    """
+def _validate_diagonal_args(array, value, xp):
+    """Validate arguments to `_fill_diagonal`/`_add_to_diagonal`."""
     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."
+        )
+    min_rows_columns = min(array.shape)
+    if value.ndim == 1 and value.shape[0] != min_rows_columns:
+        raise ValueError(
+            "`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]}"
+        )
+
+    return value, min_rows_columns
+
+
+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)
 
-    array_flat = xp.reshape(array, (-1,))
-    if add_value:
-        array_flat[:end:step] += value
+    if _is_numpy_namespace(xp):
+        xp.fill_diagonal(array, value, wrap=False)
     else:
-        array_flat[:end:step] = value
-    # `array_flat` is not always a view on `array` (e.g. for certain array types that
-    # were filled via parallel processing i.e., in `_parallel_pairwise`), thus we need
-    # to return reshaped `array_flat`.
-    return xp.reshape(array_flat, array.shape)
+        # 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
+        else:
+            for i in range(min_rows_columns):
+                array[i, i] = value[i]
+
+
+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).
+
+    Note `array` is altered in place.
+    """
+    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
+
+    # 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]
 
 
 def _is_xp_namespace(xp, name):
diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py
@@ -9,13 +9,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,
@@ -26,6 +27,7 @@
     _nanmean,
     _nanmin,
     _ravel,
+    _validate_diagonal_args,
     device,
     get_namespace,
     get_namespace_and_device,
@@ -576,21 +578,105 @@ def test_count_nonzero(
         assert device(array_xp) == device(result)
 
 
+@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("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
+
+    if function == "fill":
+        func = _fill_diagonal
+    else:
+        func = _add_to_diagonal
+
+    func(array, 1, xp)
+    assert_allclose(array.diagonal(), numpy.ones((3,)))
+
+    func(array, [0, 1, 2], xp)
+    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)
+    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", ["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, 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((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):
+        _fill_diagonal(array_xp, value=1, xp=xp)
+
+    assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np)
+
+
 @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):
+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)
 
