WIP: Enable array api support in neighbor

doc/sources/array_api.rst

@@ -96,6 +96,10 @@ The following patched classes have support for array API inputs:
 - :obj:`sklearn.linear_model.Ridge`
 - :obj:`sklearnex.linear_model.IncrementalLinearRegression`
 - :obj:`sklearnex.linear_model.IncrementalRidge`
+- :obj:`sklearn.neighbors.KNeighborsClassifier`
+- :obj:`sklearn.neighbors.KNeighborsRegressor`
+- :obj:`sklearn.neighbors.NearestNeighbors`
+- :obj:`sklearn.neighbors.LocalOutlierFactor`
 
 .. note::
     While full array API support is currently not implemented for all classes, :external+dpnp:doc:`dpnp.ndarray <reference/ndarray>`

onedal/neighbors/tests/test_knn_classification.py

@@ -19,15 +19,19 @@
 from numpy.testing import assert_array_equal
 from sklearn import datasets
 
-from onedal.neighbors import KNeighborsClassifier
 from onedal.tests.utils._device_selection import get_queues
 
+# Classification processing now happens in sklearnex layer
+from sklearnex.neighbors import KNeighborsClassifier
+
 
 @pytest.mark.parametrize("queue", get_queues())
 def test_iris(queue):
+    # queue parameter not used with sklearnex, but kept for test parametrization
     iris = datasets.load_iris()
-    clf = KNeighborsClassifier(2).fit(iris.data, iris.target, queue=queue)
-    assert clf.score(iris.data, iris.target, queue=queue) > 0.9
+    clf = KNeighborsClassifier(2).fit(iris.data, iris.target)
+    score = clf.score(iris.data, iris.target)
+    assert score > 0.9
     assert_array_equal(clf.classes_, np.sort(clf.classes_))
 
 
@@ -36,14 +40,13 @@ def test_pickle(queue):
     if queue and queue.sycl_device.is_gpu:
         pytest.skip("KNN classifier pickling for the GPU sycl_queue is buggy.")
     iris = datasets.load_iris()
-    clf = KNeighborsClassifier(2).fit(iris.data, iris.target, queue=queue)
-    expected = clf.predict(iris.data, queue=queue)
-
+    clf = KNeighborsClassifier(2).fit(iris.data, iris.target)
+    expected = clf.predict(iris.data)
     import pickle
 
     dump = pickle.dumps(clf)
     clf2 = pickle.loads(dump)
 
     assert type(clf2) == clf.__class__
-    result = clf2.predict(iris.data, queue=queue)
+    result = clf2.predict(iris.data)
     assert_array_equal(expected, result)

sklearnex/neighbors/_lof.py

@@ -29,7 +29,7 @@
 from sklearnex.neighbors.knn_unsupervised import NearestNeighbors
 
 from ..utils._array_api import get_namespace
-from ..utils.validation import check_feature_names
+from ..utils.validation import check_feature_names, validate_data
 
 
 @control_n_jobs(decorated_methods=["fit", "kneighbors", "_kneighbors"])
@@ -56,6 +56,7 @@ def _onedal_fit(self, X, y, queue=None):
         if sklearn_check_version("1.2"):
             self._validate_params()
 
+        # Let _onedal_knn_fit (NearestNeighbors._onedal_fit) handle validation
         self._onedal_knn_fit(X, y, queue=queue)
 
         if self.contamination != "auto":
@@ -74,7 +75,6 @@ def _onedal_fit(self, X, y, queue=None):
                 % (self.n_neighbors, n_samples)
             )
         self.n_neighbors_ = max(1, min(self.n_neighbors, n_samples - 1))
-
         (
             self._distances_fit_X_,
             _neighbors_indices_fit_X_,
@@ -108,11 +108,10 @@ def _onedal_fit(self, X, y, queue=None):
                     "Duplicate values are leading to incorrect results. "
                     "Increase the number of neighbors for more accurate results."
                 )
-
         return self
 
     def fit(self, X, y=None):
-        result = dispatch(
+        return dispatch(
             self,
             "fit",
             {
@@ -122,7 +121,6 @@ def fit(self, X, y=None):
             X,
             None,
         )
-        return result
 
     def _predict(self, X=None):
         check_is_fitted(self)
@@ -135,7 +133,6 @@ def _predict(self, X=None):
         else:
             is_inlier = np.ones(self.n_samples_fit_, dtype=int)
             is_inlier[self.negative_outlier_factor_ < self.offset_] = -1
-
         return is_inlier
 
     # This had to be done because predict loses the queue when no
@@ -149,9 +146,15 @@ def fit_predict(self, X, y=None):
         return self.fit(X)._predict()
 
     def _kneighbors(self, X=None, n_neighbors=None, return_distance=True):
+        # Validate n_neighbors parameter first
+        if n_neighbors is not None:
+            self._validate_n_neighbors(n_neighbors)
+
         check_is_fitted(self)
-        if X is not None:
-            check_feature_names(self, X, reset=False)
+
+        # Validate kneighbors parameters (inherited from KNeighborsDispatchingBase)
+        self._kneighbors_validation(X, n_neighbors)
+
         return dispatch(
             self,
             "kneighbors",
@@ -172,6 +175,16 @@ def _kneighbors(self, X=None, n_neighbors=None, return_distance=True):
     def score_samples(self, X):
         check_is_fitted(self)
 
+        # Validate and convert X
+        xp, _ = get_namespace(X)
+        X = validate_data(
+            self,
+            X,
+            dtype=[xp.float64, xp.float32],
+            accept_sparse="csr",
+            reset=False,
+        )
+
         distances_X, neighbors_indices_X = self._kneighbors(
             X, n_neighbors=self.n_neighbors_
         )