scikit-learn-intelex 2024.1.0__py311-none-manylinux1_x86_64.whl → 2024.2.0__py311-none-manylinux1_x86_64.whl

sklearnex/preview/cluster/k_means.py

@@ -29,18 +29,15 @@ if daal_check_version((2023, "P", 200)):
         check_is_fitted,
     )
 
-    from daal4py.sklearn._utils import (
-        control_n_jobs,
-        run_with_n_jobs,
-        sklearn_check_version,
-    )
+    from daal4py.sklearn._n_jobs_support import control_n_jobs
+    from daal4py.sklearn._utils import sklearn_check_version
     from onedal.cluster import KMeans as onedal_KMeans
 
     from ..._device_offload import dispatch, wrap_output_data
     from ..._utils import PatchingConditionsChain
     from ._common import BaseKMeans
 
-    @control_n_jobs
+    @control_n_jobs(decorated_methods=["fit", "predict"])
     class KMeans(sklearn_KMeans, BaseKMeans):
         __doc__ = sklearn_KMeans.__doc__
         n_iter_, inertia_ = None, None
@@ -212,7 +209,6 @@ if daal_check_version((2023, "P", 200)):
 
             return self
 
-        @run_with_n_jobs
         def _onedal_fit(self, X, _, sample_weight, queue=None):
             assert sample_weight is None
 
@@ -294,7 +290,6 @@ if daal_check_version((2023, "P", 200)):
                 X,
             )
 
-        @run_with_n_jobs
         def _onedal_predict(self, X, queue=None):
             X = self._validate_data(
                 X, accept_sparse=False, reset=False, dtype=[np.float64, np.float32]

sklearnex/preview/covariance/covariance.py

@@ -14,32 +14,47 @@
 # limitations under the License.
 # ===============================================================================
 
+import warnings
+
+import numpy as np
 from scipy import sparse as sp
 from sklearn.covariance import EmpiricalCovariance as sklearn_EmpiricalCovariance
 from sklearn.utils import check_array
 
-from daal4py.sklearn._utils import control_n_jobs, run_with_n_jobs, sklearn_check_version
+from daal4py.sklearn._n_jobs_support import control_n_jobs
+from daal4py.sklearn._utils import sklearn_check_version
 from onedal.common.hyperparameters import get_hyperparameters
 from onedal.covariance import EmpiricalCovariance as onedal_EmpiricalCovariance
+from sklearnex import config_context
+from sklearnex.metrics import pairwise_distances
 
-from ..._device_offload import dispatch
+from ..._device_offload import dispatch, wrap_output_data
 from ..._utils import PatchingConditionsChain, register_hyperparameters
 
 
 @register_hyperparameters({"fit": get_hyperparameters("covariance", "compute")})
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "mahalanobis"])
 class EmpiricalCovariance(sklearn_EmpiricalCovariance):
     __doc__ = sklearn_EmpiricalCovariance.__doc__
 
+    if sklearn_check_version("1.2"):
+        _parameter_constraints: dict = {
+            **sklearn_EmpiricalCovariance._parameter_constraints,
+        }
+
     def _save_attributes(self):
         assert hasattr(self, "_onedal_estimator")
-        self.covariance_ = self._onedal_estimator.covariance_
+        self._set_covariance(self._onedal_estimator.covariance_)
         self.location_ = self._onedal_estimator.location_
 
     _onedal_covariance = staticmethod(onedal_EmpiricalCovariance)
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, queue=None):
+        if X.shape[0] == 1:
+            warnings.warn(
+                "Only one sample available. You may want to reshape your data array"
+            )
+
         onedal_params = {
             "method": "dense",
             "bias": True,
@@ -54,7 +69,7 @@ class EmpiricalCovariance(sklearn_EmpiricalCovariance):
         patching_status = PatchingConditionsChain(
             f"sklearn.covariance.{class_name}.{method_name}"
         )
-        if method_name == "fit":
+        if method_name in ["fit", "mahalanobis"]:
             (X,) = data
             patching_status.and_conditions(
                 [
@@ -62,10 +77,6 @@ class EmpiricalCovariance(sklearn_EmpiricalCovariance):
                         self.assume_centered == False,
                         "assume_centered parameter is not supported on oneDAL side",
                     ),
-                    (
-                        self.store_precision == False,
-                        "precision matrix calculation is not supported on oneDAL side",
-                    ),
                     (not sp.issparse(X), "X is sparse. Sparse input is not supported."),
                 ]
             )
@@ -79,9 +90,9 @@ class EmpiricalCovariance(sklearn_EmpiricalCovariance):
         if sklearn_check_version("1.2"):
             self._validate_params()
         if sklearn_check_version("0.23"):
-            self._validate_data(X)
+            X = self._validate_data(X, force_all_finite=False)
         else:
-            check_array(X)
+            X = check_array(X, force_all_finite=False)
 
         dispatch(
             self,
@@ -95,4 +106,27 @@ class EmpiricalCovariance(sklearn_EmpiricalCovariance):
 
         return self
 
+    # expose sklearnex pairwise_distances if mahalanobis distance eventually supported
+    @wrap_output_data
+    def mahalanobis(self, X):
+        if sklearn_check_version("1.0"):
+            X = self._validate_data(X, reset=False)
+        else:
+            X = check_array(X)
+
+        precision = self.get_precision()
+        with config_context(assume_finite=True):
+            # compute mahalanobis distances
+            dist = pairwise_distances(
+                X, self.location_[np.newaxis, :], metric="mahalanobis", VI=precision
+            )
+
+        return np.reshape(dist, (len(X),)) ** 2
+
+    error_norm = wrap_output_data(sklearn_EmpiricalCovariance.error_norm)
+    score = wrap_output_data(sklearn_EmpiricalCovariance.score)
+
     fit.__doc__ = sklearn_EmpiricalCovariance.fit.__doc__
+    mahalanobis.__doc__ = sklearn_EmpiricalCovariance.mahalanobis
+    error_norm.__doc__ = sklearn_EmpiricalCovariance.error_norm.__doc__
+    score.__doc__ = sklearn_EmpiricalCovariance.score.__doc__

sklearnex/preview/decomposition/pca.py

@@ -23,7 +23,8 @@ from sklearn.base import BaseEstimator
 from sklearn.utils.extmath import stable_cumsum
 from sklearn.utils.validation import check_array, check_is_fitted
 
-from daal4py.sklearn._utils import control_n_jobs, run_with_n_jobs, sklearn_check_version
+from daal4py.sklearn._n_jobs_support import control_n_jobs
+from daal4py.sklearn._utils import sklearn_check_version
 from onedal.utils import _check_array
 
 from ..._device_offload import dispatch
@@ -43,7 +44,7 @@ from onedal.decomposition import PCA as onedal_PCA
 
 
 @register_hyperparameters({"fit": get_hyperparameters("covariance", "compute")})
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "transform"])
 class PCA(sklearn_PCA):
     __doc__ = sklearn_PCA.__doc__
 
@@ -220,7 +221,6 @@ class PCA(sklearn_PCA):
     def _onedal_gpu_supported(self, method_name, *data):
         return self._onedal_supported(method_name, *data)
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, y=None, queue=None):
         if self.n_components == "mle" or self.n_components is None:
             onedal_n_components = min(X.shape)
@@ -244,11 +244,9 @@ class PCA(sklearn_PCA):
 
         return U, S, V
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
         return self._onedal_estimator.predict(X, queue)
 
-    @run_with_n_jobs
     def _onedal_transform(self, X):
         X = _check_array(X, dtype=[np.float64, np.float32], ensure_2d=True, copy=False)
 

sklearnex/spmd/__init__.py

@@ -17,6 +17,7 @@
 __all__ = [
     "basic_statistics",
     "cluster",
+    "covariance",
     "decomposition",
     "ensemble",
     "linear_model",

sklearnex/spmd/covariance/__init__.py

@@ -0,0 +1,19 @@
+# ==============================================================================
+# 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 .covariance import EmpiricalCovariance
+
+__all__ = ["EmpiricalCovariance"]

sklearnex/spmd/covariance/covariance.py

@@ -0,0 +1,21 @@
+# ==============================================================================
+# 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 onedal.spmd.covariance import EmpiricalCovariance
+
+# TODO:
+# Currently it uses `onedal` module interface.
+# Add sklearnex dispatching.

sklearnex/spmd/linear_model/__init__.py

@@ -15,5 +15,6 @@
 # ==============================================================================
 
 from .linear_model import LinearRegression
+from .logistic_regression import LogisticRegression
 
-__all__ = ["LinearRegression"]
+__all__ = ["LinearRegression", "LogisticRegression"]

sklearnex/spmd/linear_model/logistic_regression.py

@@ -0,0 +1,21 @@
+# ==============================================================================
+# 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 onedal.spmd.linear_model import LogisticRegression
+
+# TODO:
+# Currently it uses `onedal` module interface.
+# Add sklearnex dispatching.

sklearnex/svm/nusvc.py

@@ -18,7 +18,8 @@ from sklearn.exceptions import NotFittedError
 from sklearn.svm import NuSVC as sklearn_NuSVC
 from sklearn.utils.validation import _deprecate_positional_args
 
-from daal4py.sklearn._utils import control_n_jobs, run_with_n_jobs, sklearn_check_version
+from daal4py.sklearn._n_jobs_support import control_n_jobs
+from daal4py.sklearn._utils import sklearn_check_version
 
 from .._device_offload import dispatch, wrap_output_data
 from ._common import BaseSVC
@@ -29,7 +30,9 @@ if sklearn_check_version("1.0"):
 from onedal.svm import NuSVC as onedal_NuSVC
 
 
-@control_n_jobs
+@control_n_jobs(
+    decorated_methods=["fit", "predict", "_predict_proba", "decision_function"]
+)
 class NuSVC(sklearn_NuSVC, BaseSVC):
     __doc__ = sklearn_NuSVC.__doc__
 
@@ -229,7 +232,6 @@ class NuSVC(sklearn_NuSVC, BaseSVC):
             X,
         )
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, y, sample_weight=None, queue=None):
         onedal_params = {
             "nu": self.nu,
@@ -253,11 +255,9 @@ class NuSVC(sklearn_NuSVC, BaseSVC):
             self._fit_proba(X, y, sample_weight, queue=queue)
         self._save_attributes()
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
         return self._onedal_estimator.predict(X, queue=queue)
 
-    @run_with_n_jobs
     def _onedal_predict_proba(self, X, queue=None):
         if getattr(self, "clf_prob", None) is None:
             raise NotFittedError(
@@ -272,6 +272,5 @@ class NuSVC(sklearn_NuSVC, BaseSVC):
         with config_context(**cfg):
             return self.clf_prob.predict_proba(X)
 
-    @run_with_n_jobs
     def _onedal_decision_function(self, X, queue=None):
         return self._onedal_estimator.decision_function(X, queue=queue)

sklearnex/svm/nusvr.py

@@ -17,14 +17,15 @@
 from sklearn.svm import NuSVR as sklearn_NuSVR
 from sklearn.utils.validation import _deprecate_positional_args
 
-from daal4py.sklearn._utils import control_n_jobs, run_with_n_jobs, sklearn_check_version
+from daal4py.sklearn._n_jobs_support import control_n_jobs
+from daal4py.sklearn._utils import sklearn_check_version
 from onedal.svm import NuSVR as onedal_NuSVR
 
 from .._device_offload import dispatch, wrap_output_data
 from ._common import BaseSVR
 
 
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "predict"])
 class NuSVR(sklearn_NuSVR, BaseSVR):
     __doc__ = sklearn_NuSVR.__doc__
 
@@ -142,7 +143,6 @@ class NuSVR(sklearn_NuSVR, BaseSVR):
             X,
         )
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, y, sample_weight=None, queue=None):
         onedal_params = {
             "C": self.C,
@@ -161,6 +161,5 @@ class NuSVR(sklearn_NuSVR, BaseSVR):
         self._onedal_estimator.fit(X, y, sample_weight, queue=queue)
         self._save_attributes()
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
         return self._onedal_estimator.predict(X, queue=queue)

sklearnex/svm/svc.py

@@ -20,7 +20,8 @@ from sklearn.exceptions import NotFittedError
 from sklearn.svm import SVC as sklearn_SVC
 from sklearn.utils.validation import _deprecate_positional_args
 
-from daal4py.sklearn._utils import control_n_jobs, run_with_n_jobs, sklearn_check_version
+from daal4py.sklearn._n_jobs_support import control_n_jobs
+from daal4py.sklearn._utils import sklearn_check_version
 
 from .._device_offload import dispatch, wrap_output_data
 from .._utils import PatchingConditionsChain
@@ -32,7 +33,9 @@ if sklearn_check_version("1.0"):
 from onedal.svm import SVC as onedal_SVC
 
 
-@control_n_jobs
+@control_n_jobs(
+    decorated_methods=["fit", "predict", "_predict_proba", "decision_function"]
+)
 class SVC(sklearn_SVC, BaseSVC):
     __doc__ = sklearn_SVC.__doc__
 
@@ -258,7 +261,6 @@ class SVC(sklearn_SVC, BaseSVC):
             return patching_status
         raise RuntimeError(f"Unknown method {method_name} in {class_name}")
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, y, sample_weight=None, queue=None):
         onedal_params = {
             "C": self.C,
@@ -282,11 +284,9 @@ class SVC(sklearn_SVC, BaseSVC):
             self._fit_proba(X, y, sample_weight, queue=queue)
         self._save_attributes()
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
         return self._onedal_estimator.predict(X, queue=queue)
 
-    @run_with_n_jobs
     def _onedal_predict_proba(self, X, queue=None):
         if getattr(self, "clf_prob", None) is None:
             raise NotFittedError(
@@ -301,6 +301,5 @@ class SVC(sklearn_SVC, BaseSVC):
         with config_context(**cfg):
             return self.clf_prob.predict_proba(X)
 
-    @run_with_n_jobs
     def _onedal_decision_function(self, X, queue=None):
         return self._onedal_estimator.decision_function(X, queue=queue)

sklearnex/svm/svr.py

@@ -17,14 +17,15 @@
 from sklearn.svm import SVR as sklearn_SVR
 from sklearn.utils.validation import _deprecate_positional_args
 
-from daal4py.sklearn._utils import control_n_jobs, run_with_n_jobs, sklearn_check_version
+from daal4py.sklearn._n_jobs_support import control_n_jobs
+from daal4py.sklearn._utils import sklearn_check_version
 from onedal.svm import SVR as onedal_SVR
 
 from .._device_offload import dispatch, wrap_output_data
 from ._common import BaseSVR
 
 
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "predict"])
 class SVR(sklearn_SVR, BaseSVR):
     __doc__ = sklearn_SVR.__doc__
 
@@ -143,7 +144,6 @@ class SVR(sklearn_SVR, BaseSVR):
             X,
         )
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, y, sample_weight=None, queue=None):
         onedal_params = {
             "C": self.C,
@@ -162,6 +162,5 @@ class SVR(sklearn_SVR, BaseSVR):
         self._onedal_estimator.fit(X, y, sample_weight, queue=queue)
         self._save_attributes()
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
         return self._onedal_estimator.predict(X, queue=queue)

sklearnex/tests/test_memory_usage.py

@@ -94,10 +94,7 @@ def remove_duplicated_estimators(estimators_list):
     return estimators_map.values()
 
 
-BANNED_ESTIMATORS = (
-    "LocalOutlierFactor",  # fails on ndarray_c for sklearn > 1.0
-    "TSNE",  # too slow for using in testing on common data size
-)
+BANNED_ESTIMATORS = ("TSNE",)  # too slow for using in testing on common data size
 estimators = [
     PreviewPCA,
     TrainTestSplitEstimator,

sklearnex/tests/test_monkeypatch.py

@@ -45,8 +45,12 @@ def test_monkey_patching():
         n = _classes[i][1]
 
         sklearnex.unpatch_sklearn(t)
-        class_module = getattr(p, n).__module__
-        assert class_module.startswith("sklearn"), "Unpatching has completed with error."
+        sklearn_class = getattr(p, n, None)
+        if sklearn_class is not None:
+            sklearn_class = sklearn_class.__module__
+        assert sklearn_class is None or sklearn_class.startswith(
+            "sklearn"
+        ), "Unpatching has completed with error."
 
     sklearnex.unpatch_sklearn()
 
@@ -55,8 +59,12 @@ def test_monkey_patching():
         p = _classes[i][0]
         n = _classes[i][1]
 
-        class_module = getattr(p, n).__module__
-        assert class_module.startswith("sklearn"), "Unpatching has completed with error."
+        sklearn_class = getattr(p, n, None)
+        if sklearn_class is not None:
+            sklearn_class = sklearn_class.__module__
+        assert sklearn_class is None or sklearn_class.startswith(
+            "sklearn"
+        ), "Unpatching has completed with error."
 
     sklearnex.unpatch_sklearn()
 
@@ -85,7 +93,10 @@ def test_patch_by_list_simple():
 
     assert RandomForestRegressor.__module__.startswith("sklearn")
     assert KNeighborsRegressor.__module__.startswith("sklearn")
-    assert LogisticRegression.__module__.startswith("daal4py")
+    if daal_check_version((2024, "P", 1)):
+        assert LogisticRegression.__module__.startswith("sklearnex")
+    else:
+        assert LogisticRegression.__module__.startswith("daal4py")
     assert SVC.__module__.startswith("sklearn")
 
     sklearnex.unpatch_sklearn()
@@ -101,7 +112,10 @@ def test_patch_by_list_many_estimators():
 
     assert RandomForestRegressor.__module__.startswith("sklearn")
     assert KNeighborsRegressor.__module__.startswith("sklearn")
-    assert LogisticRegression.__module__.startswith("daal4py")
+    if daal_check_version((2024, "P", 1)):
+        assert LogisticRegression.__module__.startswith("sklearnex")
+    else:
+        assert LogisticRegression.__module__.startswith("daal4py")
     assert SVC.__module__.startswith("daal4py") or SVC.__module__.startswith("sklearnex")
 
     sklearnex.unpatch_sklearn()
@@ -119,7 +133,10 @@ def test_unpatch_by_list_many_estimators():
     assert KNeighborsRegressor.__module__.startswith(
         "daal4py"
     ) or KNeighborsRegressor.__module__.startswith("sklearnex")
-    assert LogisticRegression.__module__.startswith("daal4py")
+    if daal_check_version((2024, "P", 1)):
+        assert LogisticRegression.__module__.startswith("sklearnex")
+    else:
+        assert LogisticRegression.__module__.startswith("daal4py")
     assert SVC.__module__.startswith("daal4py") or SVC.__module__.startswith("sklearnex")
 
     sklearnex.unpatch_sklearn(["KNeighborsRegressor", "RandomForestRegressor"])
@@ -131,7 +148,11 @@ def test_unpatch_by_list_many_estimators():
 
     assert RandomForestRegressor.__module__.startswith("sklearn")
     assert KNeighborsRegressor.__module__.startswith("sklearn")
-    assert LogisticRegression.__module__.startswith("daal4py")
+    if daal_check_version((2024, "P", 1)):
+        assert LogisticRegression.__module__.startswith("sklearnex")
+    else:
+        assert LogisticRegression.__module__.startswith("daal4py")
+
     assert SVC.__module__.startswith("daal4py") or SVC.__module__.startswith("sklearnex")
 
 
@@ -161,12 +182,11 @@ def test_preview_namespace():
         from sklearn.cluster import DBSCAN
         from sklearn.decomposition import PCA
         from sklearn.ensemble import RandomForestClassifier
-        from sklearn.linear_model import LinearRegression, LogisticRegression
+        from sklearn.linear_model import LinearRegression
         from sklearn.svm import SVC
 
         return (
             LinearRegression(),
-            LogisticRegression(),
             PCA(),
             DBSCAN(),
             SVC(),
@@ -182,7 +202,7 @@ def test_preview_namespace():
 
     assert _is_preview_enabled()
 
-    lr, log_reg, pca, dbscan, svc, rfc = get_estimators()
+    lr, pca, dbscan, svc, rfc = get_estimators()
     assert "sklearnex" in rfc.__module__
 
     if daal_check_version((2023, "P", 100)):
@@ -190,20 +210,14 @@ def test_preview_namespace():
     else:
         assert "daal4py" in lr.__module__
 
-    if daal_check_version((2024, "P", 1)):
-        assert "sklearnex" in log_reg.__module__
-    else:
-        assert "daal4py" in log_reg.__module__
-
     assert "sklearnex.preview" in pca.__module__
     assert "sklearnex" in dbscan.__module__
     assert "sklearnex" in svc.__module__
     sklearnex.unpatch_sklearn()
 
     # no patching behavior
-    lr, log_reg, pca, dbscan, svc, rfc = get_estimators()
+    lr, pca, dbscan, svc, rfc = get_estimators()
     assert "sklearn." in lr.__module__ and "daal4py" not in lr.__module__
-    assert "sklearn." in log_reg.__module__ and "daal4py" not in log_reg.__module__
     assert "sklearn." in pca.__module__ and "daal4py" not in pca.__module__
     assert "sklearn." in dbscan.__module__ and "daal4py" not in dbscan.__module__
     assert "sklearn." in svc.__module__ and "daal4py" not in svc.__module__
@@ -213,13 +227,12 @@ def test_preview_namespace():
     sklearnex.patch_sklearn()
     assert not _is_preview_enabled()
 
-    lr, log_reg, pca, dbscan, svc, rfc = get_estimators()
+    lr, pca, dbscan, svc, rfc = get_estimators()
     if daal_check_version((2023, "P", 100)):
         assert "sklearnex" in lr.__module__
     else:
         assert "daal4py" in lr.__module__
 
-    assert "daal4py" in log_reg.__module__
     assert "daal4py" in pca.__module__
     assert "sklearnex" in rfc.__module__
     assert "sklearnex" in dbscan.__module__

sklearnex/tests/test_n_jobs_support.py

@@ -14,18 +14,80 @@
 # limitations under the License.
 # ==============================================================================
 
+import inspect
+import logging
+from multiprocessing import cpu_count
+
 import pytest
+from sklearn.base import BaseEstimator
+from sklearn.datasets import make_classification
+
+from sklearnex.dispatcher import get_patch_map
+from sklearnex.svm import SVC, NuSVC
+
+ESTIMATORS = set(
+    filter(
+        lambda x: inspect.isclass(x) and issubclass(x, BaseEstimator),
+        [value[0][0][2] for value in get_patch_map().values()],
+    )
+)
+
+X, Y = make_classification(n_samples=40, n_features=4, random_state=42)
+
+
+@pytest.mark.parametrize("estimator_class", ESTIMATORS)
+@pytest.mark.parametrize("n_jobs", [None, -1, 1, 2])
+def test_n_jobs_support(caplog, estimator_class, n_jobs):
+    def check_estimator_doc(estimator):
+        if estimator.__doc__ is not None:
+            assert "n_jobs" in estimator.__doc__
+
+    def check_n_jobs_entry_in_logs(caplog, function_name, n_jobs):
+        for rec in caplog.records:
+            if function_name in rec.message and "threads" in rec.message:
+                expected_n_jobs = n_jobs if n_jobs > 0 else cpu_count() + 1 + n_jobs
+                logging.info(f"{function_name}: setting {expected_n_jobs} threads")
+                if f"{function_name}: setting {expected_n_jobs} threads" in rec.message:
+                    return True
+        # False if n_jobs is set and not found in logs
+        return n_jobs is None
 
-from sklearnex.cluster import KMeans
-from sklearnex.linear_model import ElasticNet, Lasso, Ridge
-from sklearnex.svm import SVC, SVR, NuSVC, NuSVR
+    def check_method(*args, method, caplog):
+        method(*args)
+        assert check_n_jobs_entry_in_logs(caplog, method.__name__, n_jobs)
 
-estimators = [KMeans, SVC, SVR, NuSVC, NuSVR, Lasso, Ridge, ElasticNet]
+    def check_methods_decoration(estimator):
+        funcs = {
+            i: getattr(estimator, i)
+            for i in dir(estimator)
+            if hasattr(estimator, i) and callable(getattr(estimator, i))
+        }
 
+        for func_name, func in funcs.items():
+            assert hasattr(func, "__onedal_n_jobs_decorated__") == (
+                func_name in estimator._n_jobs_supported_onedal_methods
+            ), f"{estimator}.{func_name} n_jobs decoration does not match {estimator} n_jobs supported methods"
 
-@pytest.mark.parametrize("estimator", estimators)
-def test_n_jobs_support(estimator):
-    # use `n_jobs` parameter where original sklearn doesn't expect it
-    estimator(n_jobs=1)
+    caplog.set_level(logging.DEBUG, logger="sklearnex")
+    estimator_kwargs = {"n_jobs": n_jobs}
+    # by default, [Nu]SVC.predict_proba is restricted by @available_if decorator
+    if estimator_class in [SVC, NuSVC]:
+        estimator_kwargs["probability"] = True
+    estimator_instance = estimator_class(**estimator_kwargs)
     # check `n_jobs` parameter doc entry
-    assert "n_jobs" in estimator.__doc__
+    check_estimator_doc(estimator_class)
+    check_estimator_doc(estimator_instance)
+    # check `n_jobs` log entry for supported methods
+    # `fit` call is required before other methods
+    check_method(X, Y, method=estimator_instance.fit, caplog=caplog)
+    for method_name in estimator_instance._n_jobs_supported_onedal_methods:
+        if method_name == "fit":
+            continue
+        method = getattr(estimator_instance, method_name)
+        if len(inspect.signature(method).parameters) == 0:
+            check_method(method=method, caplog=caplog)
+        else:
+            check_method(X, method=method, caplog=caplog)
+    # check if correct methods were decorated
+    check_methods_decoration(estimator_class)
+    check_methods_decoration(estimator_instance)