fix: move more parts to sklearnex

Author: yuejiaointel

sklearnex/neighbors/common.py

@@ -250,6 +250,32 @@ def _compute_weights(self, distances, weights_param):
                 "'distance', or a callable function"
             )
 
+    def _process_self_neighbors(self, results, X, n_neighbors, return_distance):
+        """Process results to remove self-neighbors for training data queries."""
+        
+        if return_distance:
+            neigh_dist, neigh_ind = results
+        else:
+            neigh_ind = results
+
+        n_queries = X.shape[0] if X is not None else self._fit_X.shape[0]
+        sample_range = np.arange(n_queries)[:, None]
+        sample_mask = neigh_ind != sample_range
+
+        # Corner case: When the number of duplicates are more
+        # than the number of neighbors, the first NN will not
+        # be the sample, but a duplicate.
+        # In that case mask the first duplicate.
+        dup_gr_nbrs = np.all(sample_mask, axis=1)
+        sample_mask[:, 0][dup_gr_nbrs] = False
+
+        neigh_ind = np.reshape(neigh_ind[sample_mask], (n_queries, n_neighbors))
+
+        if return_distance:
+            neigh_dist = np.reshape(neigh_dist[sample_mask], (n_queries, n_neighbors))
+            return neigh_dist, neigh_ind
+        return neigh_ind
+
     def _validate_feature_count(self, X, method_name=""):
         n_features = getattr(self, "n_features_in_", None)
         shape = getattr(X, "shape", None)

sklearnex/neighbors/knn_classification.py

@@ -175,14 +175,13 @@ def score(self, X, y, sample_weight=None):
     def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         check_is_fitted(self)
 
-        # Handle X=None case at sklearnex level
-        if X is None:
-            query_is_train = True
+        # Handle X=None case and determine if we need self-neighbor processing at sklearnex level
+        query_is_train = X is None
+        if query_is_train:
             X = self._fit_X
-            # Include an extra neighbor to account for the sample itself being returned
+            # For training data queries, we need to get extra neighbors to remove self-matches
             effective_n_neighbors = (n_neighbors if n_neighbors is not None else self.n_neighbors) + 1
         else:
-            query_is_train = False
             effective_n_neighbors = n_neighbors if n_neighbors is not None else self.n_neighbors
             check_feature_names(self, X, reset=False)
             # Perform preprocessing at sklearnex level
@@ -193,7 +192,7 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         # Validate kneighbors parameters
         self._validate_kneighbors_params(n_neighbors, X)
 
-        # Call oneDAL with the effective n_neighbors
+        # Call oneDAL with the effective n_neighbors (all preprocessing done)
         result = dispatch(
             self,
             "kneighbors",
@@ -204,14 +203,14 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
                 "sklearn": _sklearn_KNeighborsClassifier.kneighbors,
             },
             X,
-            n_neighbors=n_neighbors,
+            n_neighbors=effective_n_neighbors,  # Pass effective_n_neighbors to sklearn too
             return_distance=return_distance,
         )
 
-        # Process results at sklearnex level
+        # If query was training data, post-process to remove self-neighbors at sklearnex level
         if query_is_train:
             final_n_neighbors = n_neighbors if n_neighbors is not None else self.n_neighbors
-            result = self._process_kneighbors_results(result, X, final_n_neighbors, return_distance, query_is_train)
+            result = self._process_self_neighbors(result, X, final_n_neighbors, return_distance)
 
         return result
 

sklearnex/neighbors/knn_regression.py

@@ -153,14 +153,13 @@ def score(self, X, y, sample_weight=None):
     def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         check_is_fitted(self)
 
-        # Handle X=None case at sklearnex level
-        if X is None:
-            query_is_train = True
+        # Handle X=None case and determine if we need self-neighbor processing at sklearnex level
+        query_is_train = X is None
+        if query_is_train:
             X = self._fit_X
-            # Include an extra neighbor to account for the sample itself being returned
+            # For training data queries, we need to get extra neighbors to remove self-matches
             effective_n_neighbors = (n_neighbors if n_neighbors is not None else self.n_neighbors) + 1
         else:
-            query_is_train = False
             effective_n_neighbors = n_neighbors if n_neighbors is not None else self.n_neighbors
             check_feature_names(self, X, reset=False)
             # Perform preprocessing at sklearnex level
@@ -171,7 +170,7 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         # Validate kneighbors parameters
         self._validate_kneighbors_params(n_neighbors, X)
 
-        # Call oneDAL with the effective n_neighbors
+        # Call oneDAL with the effective n_neighbors (all preprocessing done)
         result = dispatch(
             self,
             "kneighbors",
@@ -182,14 +181,14 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
                 "sklearn": _sklearn_KNeighborsRegressor.kneighbors,
             },
             X,
-            n_neighbors=n_neighbors,
+            n_neighbors=effective_n_neighbors,  # Pass effective_n_neighbors to sklearn too
             return_distance=return_distance,
         )
 
-        # Process results at sklearnex level
+        # If query was training data, post-process to remove self-neighbors at sklearnex level
         if query_is_train:
             final_n_neighbors = n_neighbors if n_neighbors is not None else self.n_neighbors
-            result = self._process_kneighbors_results(result, X, final_n_neighbors, return_distance, query_is_train)
+            result = self._process_self_neighbors(result, X, final_n_neighbors, return_distance)
 
         return result
 

sklearnex/neighbors/knn_unsupervised.py

@@ -76,14 +76,13 @@ def fit(self, X, y=None):
     def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         check_is_fitted(self)
 
-        # Handle X=None case at sklearnex level
-        if X is None:
-            query_is_train = True
+        # Handle X=None case and determine if we need self-neighbor processing at sklearnex level
+        query_is_train = X is None
+        if query_is_train:
             X = self._fit_X
-            # Include an extra neighbor to account for the sample itself being returned
+            # For training data queries, we need to get extra neighbors to remove self-matches
             effective_n_neighbors = (n_neighbors if n_neighbors is not None else self.n_neighbors) + 1
         else:
-            query_is_train = False
             effective_n_neighbors = n_neighbors if n_neighbors is not None else self.n_neighbors
             check_feature_names(self, X, reset=False)
             # Perform preprocessing at sklearnex level
@@ -94,7 +93,7 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         # Validate kneighbors parameters
         self._validate_kneighbors_params(n_neighbors, X)
 
-        # Call oneDAL with the effective n_neighbors
+        # Call oneDAL with the effective n_neighbors (all preprocessing done)
         result = dispatch(
             self,
             "kneighbors",
@@ -105,14 +104,14 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
                 "sklearn": _sklearn_NearestNeighbors.kneighbors,
             },
             X,
-            n_neighbors=n_neighbors,
+            n_neighbors=effective_n_neighbors,  # Pass effective_n_neighbors to sklearn too
             return_distance=return_distance,
         )
 
-        # Process results at sklearnex level
+        # If query was training data, post-process to remove self-neighbors at sklearnex level
         if query_is_train:
             final_n_neighbors = n_neighbors if n_neighbors is not None else self.n_neighbors
-            result = self._process_kneighbors_results(result, X, final_n_neighbors, return_distance, query_is_train)
+            result = self._process_self_neighbors(result, X, final_n_neighbors, return_distance)
 
         return result