MAINT: adding common BaseEstimatorSPMD for spmd ifaces (#1679)

onedal/spmd/_common.py

@@ -0,0 +1,24 @@
+# ==============================================================================
+# Copyright 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from abc import ABC
+
+from ..common._spmd_policy import _get_spmd_policy
+
+
+class BaseEstimatorSPMD(ABC):
+    def _get_policy(self, queue, *data):
+        return _get_spmd_policy(queue)

onedal/spmd/basic_statistics/basic_statistics.py

@@ -14,20 +14,13 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import ABC
-
 from onedal.basic_statistics import BasicStatistics as BasicStatistics_Batch
 
 from ..._device_offload import support_usm_ndarray
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class BaseBasicStatisticsSPMD(ABC):
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class BasicStatistics(BaseBasicStatisticsSPMD, BasicStatistics_Batch):
+class BasicStatistics(BaseEstimatorSPMD, BasicStatistics_Batch):
     @support_usm_ndarray()
     def compute(self, data, weights=None, queue=None):
         return super().compute(data, weights, queue)

onedal/spmd/cluster/dbscan.py

@@ -14,17 +14,10 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import ABC
-
 from onedal.cluster import DBSCAN as DBSCAN_Batch
 
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class BaseDBSCANspmd(ABC):
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class DBSCAN(BaseDBSCANspmd, DBSCAN_Batch):
+class DBSCAN(BaseEstimatorSPMD, DBSCAN_Batch):
     pass

onedal/spmd/cluster/kmeans.py

@@ -14,20 +14,13 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import ABC
-
 from onedal.cluster import KMeans as KMeans_Batch
 
 from ..._device_offload import support_usm_ndarray
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class BaseKMeansSPMD(ABC):
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class KMeans(BaseKMeansSPMD, KMeans_Batch):
+class KMeans(BaseEstimatorSPMD, KMeans_Batch):
     @support_usm_ndarray()
     def fit(self, X, queue=None):
         return super().fit(X, queue)

onedal/spmd/decomposition/pca.py

@@ -14,19 +14,13 @@
 # limitations under the License.
 # ==============================================================================
 
-
 from onedal.decomposition.pca import PCA as PCABatch
 
 from ..._device_offload import support_usm_ndarray
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class BasePCASPMD:
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class PCA(BasePCASPMD, PCABatch):
+class PCA(BaseEstimatorSPMD, PCABatch):
     @support_usm_ndarray()
     def fit(self, X, queue):
         return super().fit(X, queue)

onedal/spmd/ensemble/forest.py

@@ -14,22 +14,15 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import ABC
-
 from onedal.ensemble import RandomForestClassifier as RandomForestClassifier_Batch
 from onedal.ensemble import RandomForestRegressor as RandomForestRegressor_Batch
 
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class BaseForestSPMD(ABC):
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class RandomForestClassifier(BaseForestSPMD, RandomForestClassifier_Batch):
+class RandomForestClassifier(BaseEstimatorSPMD, RandomForestClassifier_Batch):
     pass
 
 
-class RandomForestRegressor(BaseForestSPMD, RandomForestRegressor_Batch):
+class RandomForestRegressor(BaseEstimatorSPMD, RandomForestRegressor_Batch):
     pass

onedal/spmd/linear_model/linear_model.py

@@ -14,20 +14,13 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import ABC
-
 from onedal.linear_model import LinearRegression as LinearRegression_Batch
 
 from ..._device_offload import support_usm_ndarray
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class BaseLinearRegressionSPMD(ABC):
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class LinearRegression(BaseLinearRegressionSPMD, LinearRegression_Batch):
+class LinearRegression(BaseEstimatorSPMD, LinearRegression_Batch):
     @support_usm_ndarray()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue)

onedal/spmd/neighbors/neighbors.py

@@ -14,21 +14,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import ABC
-
 from onedal.neighbors import KNeighborsClassifier as KNeighborsClassifier_Batch
 from onedal.neighbors import KNeighborsRegressor as KNeighborsRegressor_Batch
 
 from ..._device_offload import support_usm_ndarray
-from ...common._spmd_policy import _get_spmd_policy
-
-
-class NeighborsCommonBaseSPMD(ABC):
-    def _get_policy(self, queue, *data):
-        return _get_spmd_policy(queue)
+from .._common import BaseEstimatorSPMD
 
 
-class KNeighborsClassifier(NeighborsCommonBaseSPMD, KNeighborsClassifier_Batch):
+class KNeighborsClassifier(BaseEstimatorSPMD, KNeighborsClassifier_Batch):
     @support_usm_ndarray()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue)
@@ -46,7 +39,7 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True, queue=None)
         return super().kneighbors(X, n_neighbors, return_distance, queue)
 
 
-class KNeighborsRegressor(NeighborsCommonBaseSPMD, KNeighborsRegressor_Batch):
+class KNeighborsRegressor(BaseEstimatorSPMD, KNeighborsRegressor_Batch):
     @support_usm_ndarray()
     def fit(self, X, y, queue=None):
         if queue is not None and queue.sycl_device.is_gpu:
@@ -72,7 +65,7 @@ def _get_onedal_params(self, X, y=None):
         return params
 
 
-class NearestNeighbors(NeighborsCommonBaseSPMD):
+class NearestNeighbors(BaseEstimatorSPMD):
     @support_usm_ndarray()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue)