Fix: Array-API - avoid failing for numpy fit + predict with sparse or array-like X (#34144)

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>
Co-authored-by: Anne Beyer <anne.beyer@mailbox.org>
Co-authored-by: Omar Salman <omar.salman@arbisoft.com>

Author: 4 people

doc/whats_new/upcoming_changes/array-api/34144.fix.rst

@@ -0,0 +1,5 @@
+- Fixed a bug where NumPy-fitted estimators could raise an error with
+  ``config_context(array_api_dispatch=True)`` when making predictions with
+  array-like or SciPy sparse inputs, or when a fitted attribute was sparse,
+  such as after calling :meth:`linear_model.LogisticRegression.sparsify`.
+  By :user:`Arthur Lacote <cakedev0>`.

sklearn/utils/_array_api.py

@@ -1135,7 +1135,22 @@ def check_same_namespace(X, estimator, *, attribute, method):
     if X_xp == a_xp and X_device == a_device:
         return
 
-    if X_xp != a_xp:
+    if _is_numpy_namespace(a_xp) and _is_numpy_namespace(X_xp):
+        # this condition is reached when either:
+        # - `X` is array-like or sparse-matrix and the estimator was fitted with numpy
+        # - `attr` is a sparse matrix and `X` is a numpy array
+        # in which case devices are different (None vs "cpu") but
+        # `check_same_namespace` should not raise
+        return
+
+    if X_device is None:
+        type_name = "sparse array" if sp.issparse(X) else "array-like"
+        msg = (
+            f"Array namespace used during fit ({a_xp.__name__}) "
+            f"is not compatible with the {type_name} input passed to {method}. "
+            f"Only the NumPy namespace is compatible with {type_name} inputs."
+        )
+    elif X_xp != a_xp:
         msg = (
             f"Array namespaces used during fit ({a_xp.__name__}) "
             f"and {method} ({X_xp.__name__}) differ."
@@ -1148,7 +1163,8 @@ def check_same_namespace(X, estimator, *, attribute, method):
         "must use the same namespace and the same device as those passed to fit(). "
         f"{msg} "
         "You can move the estimator to the same namespace and device as X with: "
-        "'from sklearn.utils._array_api import move_estimator_to; "
+        "'from sklearn.utils._array_api import get_namespace_and_device, "
+        "move_estimator_to; "
         "xp, _, device = get_namespace_and_device(X); "
         "estimator = move_estimator_to(estimator, xp, device)'"
     )

sklearn/utils/tests/test_array_api.py

@@ -503,6 +503,9 @@ def predict(self, X):
         check_same_namespace(X, self, attribute="X_", method="predict")
         return X
 
+    def sparsify(self):
+        self.X_ = sp.csr_matrix(self.X_)
+
 
 class SimpleEstimatorCustomLogic(BaseEstimator):
     def fit(self, X, y=None):
@@ -607,6 +610,31 @@ def test_check_fitted_attribute():
             est.predict(numpy.asarray([0]))
 
 
+@skip_if_array_api_compat_not_configured
+@pytest.mark.parametrize("X", [[[1.3, 4.5]], sp.csr_array([[1.3, 4.5]])])
+def test_check_fitted_attribute_with_non_array_input(X):
+    """Check validation of non-array input against fitted attribute ``X_``.
+
+    ``SimpleEstimator.predict`` calls ``check_same_namespace`` with
+    ``attribute="X_"`` to compare the input with the fitted data.
+    """
+    xp = pytest.importorskip("array_api_strict")
+
+    with config_context(array_api_dispatch=True):
+        est = SimpleEstimator().fit(numpy.asarray([[1.3, 4.5]]))
+        # shouldn't raise:
+        est.predict(X)
+
+        est.sparsify()
+        # shouldn't raise either:
+        est.predict(X)
+        est.predict(numpy.asarray([[1.3, 4.5]]))
+
+        est = SimpleEstimator().fit(xp.asarray([[1.3, 4.5]]))
+        with pytest.raises(ValueError, match="Array namespace.*not compatible"):
+            est.predict(X)
+
+
 @pytest.mark.parametrize(
     "namespace, device_name, dtype_name",
     yield_namespace_device_dtype_combinations(),