Fix all NaN slice handling

betatim · betatim · commit beb130a9ad70 · 2023-08-04T16:43:43.000+02:00
When using the axis keyword it can happen that a slice contains only
NaNs. This change corrects the logic that restores the NaNs at the end.
Added tests for 2D inputs.
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -91,6 +91,7 @@ Estimators with support for `Array API`-compatible inputs
 - :class:`decomposition.PCA` (with `svd_solver="full"`,
   `svd_solver="randomized"` and `power_iteration_normalizer="QR"`)
 - :class:`discriminant_analysis.LinearDiscriminantAnalysis` (with `solver="svd"`)
+- :class:`preprocessing.MinMaxScaler`
 
 Coverage for more estimators is expected to grow over time. Please follow the
 dedicated `meta-issue on GitHub
diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
@@ -495,9 +495,13 @@ def _nanmin(X, axis=None):
         return xp.asarray(numpy.nanmin(X, axis=axis))
 
     else:
-        X = xp.min(xp.where(~xp.isnan(X), X, xp.asarray(+xp.inf)), axis=axis)
+        mask = xp.isnan(X)
+        X = xp.min(xp.where(mask, xp.asarray(+xp.inf), X), axis=axis)
         # Replace Infs from all NaN slices with NaN again
-        return xp.where(~xp.isinf(X), X, xp.asarray(xp.nan))
+        mask = xp.all(mask, axis=axis)
+        if xp.any(mask):
+            X = xp.where(~mask, X, xp.asarray(xp.nan))
+        return X
 
 
 def _nanmax(X, axis=None):
@@ -508,9 +512,13 @@ def _nanmax(X, axis=None):
         return xp.asarray(numpy.nanmax(X, axis=axis))
 
     else:
-        X = xp.max(xp.where(~xp.isnan(X), X, xp.asarray(-xp.inf)), axis=axis)
+        mask = xp.isnan(X)
+        X = xp.max(xp.where(mask, xp.asarray(-xp.inf), X), axis=axis)
         # Replace Infs from all NaN slices with NaN again
-        return xp.where(~_isneginf(X), X, xp.asarray(xp.nan))
+        mask = xp.all(mask, axis=axis)
+        if xp.any(mask):
+            X = xp.where(~mask, X, xp.asarray(xp.nan))
+        return X
 
 
 def _asarray_with_order(array, dtype=None, order=None, copy=None, *, xp=None):
diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py
@@ -1,3 +1,5 @@
+from functools import partial
+
 import numpy
 import pytest
 from numpy.testing import assert_allclose, assert_array_equal
@@ -171,17 +173,43 @@ def test_asarray_with_order_ignored():
     [
         ([1, 2, numpy.nan], _nanmin, 1),
         ([1, -2, -numpy.nan], _nanmin, -2),
+        ([numpy.inf, numpy.inf], _nanmin, numpy.inf),
+        (
+            [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]],
+            partial(_nanmin, axis=0),
+            [1.0, 2.0, 3.0],
+        ),
+        (
+            [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]],
+            partial(_nanmin, axis=1),
+            [1.0, numpy.nan, 4.0],
+        ),
         ([1, 2, numpy.nan], _nanmax, 2),
         ([1, 2, numpy.nan], _nanmax, 2),
+        ([-numpy.inf, -numpy.inf], _nanmax, -numpy.inf),
+        (
+            [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]],
+            partial(_nanmax, axis=0),
+            [4.0, 5.0, 6.0],
+        ),
+        (
+            [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]],
+            partial(_nanmax, axis=1),
+            [3.0, numpy.nan, 6.0],
+        ),
     ],
 )
 def test_nan_reductions(library, X, reduction, expected):
     """Check NaN reductions like _nanmin and _nanmax"""
     xp = pytest.importorskip(library)
 
+    if isinstance(expected, list):
+        expected = xp.asarray(expected)
+
     with config_context(array_api_dispatch=True):
         result = reduction(xp.asarray(X))
-    assert result == expected
+
+    assert_allclose(result, expected)
 
 
 @skip_if_array_api_compat_not_configured
