fix: fix shape

Author: yuejiaointel

onedal/neighbors/neighbors.py

@@ -70,6 +70,13 @@ def _get_onedal_params(self, X, y=None, n_neighbors=None):
 
             fptype = np.float64
 
+        # _fit_method should be set by sklearnex level before calling oneDAL
+        if not hasattr(self, "_fit_method") or self._fit_method is None:
+            raise ValueError(
+                "_fit_method must be set by sklearnex level before calling oneDAL. "
+                "This indicates improper usage - oneDAL neighbors should not be called directly."
+            )
+
         return {
             "fptype": fptype,
             "vote_weights": "uniform" if weights == "uniform" else "distance",

sklearnex/neighbors/common.py

@@ -166,6 +166,41 @@ def _validate_kneighbors_bounds(self, n_neighbors, query_is_train, X):
                 f"n_samples = {X.shape[0]}"  # include n_samples for common tests
             )
 
+    def _process_classification_targets(self, y):
+        """Process classification targets and set class-related attributes."""
+        import numpy as np
+
+        # Handle shape processing
+        shape = getattr(y, "shape", None)
+        self._shape = shape if shape is not None else y.shape
+
+        if y.ndim == 1 or y.ndim == 2 and y.shape[1] == 1:
+            self.outputs_2d_ = False
+            y = y.reshape((-1, 1))
+        else:
+            self.outputs_2d_ = True
+
+        # Process classes
+        self.classes_ = []
+        self._y = np.empty(y.shape, dtype=int)
+        for k in range(self._y.shape[1]):
+            classes, self._y[:, k] = np.unique(y[:, k], return_inverse=True)
+            self.classes_.append(classes)
+
+        if not self.outputs_2d_:
+            self.classes_ = self.classes_[0]
+            self._y = self._y.ravel()
+
+        return y
+
+    def _process_regression_targets(self, y):
+        """Process regression targets and set shape-related attributes."""
+        # Handle shape processing for regression
+        shape = getattr(y, "shape", None)
+        self._shape = shape if shape is not None else y.shape
+        self._y = y
+        return y
+
     def _fit_validation(self, X, y=None):
         if sklearn_check_version("1.2"):
             self._validate_params()

sklearnex/neighbors/knn_classification.py

@@ -191,6 +191,9 @@ def _onedal_fit(self, X, y, queue=None):
 
         _check_classification_targets(y)
 
+        # Handle shape and class processing at sklearnex level
+        y = self._process_classification_targets(y)
+
         onedal_params = {
             "n_neighbors": self.n_neighbors,
             "weights": self.weights,
@@ -204,6 +207,13 @@ def _onedal_fit(self, X, y, queue=None):
         self._onedal_estimator.effective_metric_ = self.effective_metric_
         self._onedal_estimator.effective_metric_params_ = self.effective_metric_params_
         self._onedal_estimator._fit_method = self._fit_method
+
+        # Set shape and class attributes on the onedal estimator
+        self._onedal_estimator._shape = self._shape
+        self._onedal_estimator.classes_ = self.classes_
+        self._onedal_estimator._y = self._y
+        self._onedal_estimator.outputs_2d_ = self.outputs_2d_
+
         self._onedal_estimator.fit(X, y, queue=queue)
 
         self._save_attributes()

sklearnex/neighbors/knn_regression.py

@@ -158,6 +158,9 @@ def _onedal_fit(self, X, y, queue=None):
         # Parse auto method
         self._fit_method = self._parse_auto_method(self.algorithm, X.shape[0], X.shape[1])
 
+        # Handle shape processing at sklearnex level
+        y = self._process_regression_targets(y)
+
         onedal_params = {
             "n_neighbors": self.n_neighbors,
             "weights": self.weights,
@@ -171,6 +174,11 @@ def _onedal_fit(self, X, y, queue=None):
         self._onedal_estimator.effective_metric_ = self.effective_metric_
         self._onedal_estimator.effective_metric_params_ = self.effective_metric_params_
         self._onedal_estimator._fit_method = self._fit_method
+
+        # Set shape attributes on the onedal estimator
+        self._onedal_estimator._shape = self._shape
+        self._onedal_estimator._y = self._y
+
         self._onedal_estimator.fit(X, y, queue=queue)
 
         self._save_attributes()

sklearnex/neighbors/knn_unsupervised.py

@@ -149,6 +149,9 @@ def _onedal_fit(self, X, y=None, queue=None):
         # Parse auto method
         self._fit_method = self._parse_auto_method(self.algorithm, X.shape[0], X.shape[1])
 
+        # Set basic attributes for unsupervised
+        self.classes_ = None
+
         onedal_params = {
             "n_neighbors": self.n_neighbors,
             "algorithm": self.algorithm,
@@ -161,6 +164,10 @@ def _onedal_fit(self, X, y=None, queue=None):
         self._onedal_estimator.effective_metric_ = self.effective_metric_
         self._onedal_estimator.effective_metric_params_ = self.effective_metric_params_
         self._onedal_estimator._fit_method = self._fit_method
+
+        # Set attributes on the onedal estimator
+        self._onedal_estimator.classes_ = self.classes_
+
         self._onedal_estimator.fit(X, y, queue=queue)
 
         self._save_attributes()