scikit-learn-intelex 2024.1.0__py311-none-win_amd64.whl → 2025.1.0__py311-none-win_amd64.whl

{scikit_learn_intelex-2024.1.0.data → scikit_learn_intelex-2025.1.0.data}/data/Lib/site-packages/sklearnex/ensemble/_forest.py

@@ -20,13 +20,16 @@ from abc import ABC
 
 import numpy as np
 from scipy import sparse as sp
-from sklearn.base import clone
-from sklearn.ensemble import ExtraTreesClassifier as sklearn_ExtraTreesClassifier
-from sklearn.ensemble import ExtraTreesRegressor as sklearn_ExtraTreesRegressor
-from sklearn.ensemble import RandomForestClassifier as sklearn_RandomForestClassifier
-from sklearn.ensemble import RandomForestRegressor as sklearn_RandomForestRegressor
+from sklearn.base import BaseEstimator, clone
+from sklearn.ensemble import ExtraTreesClassifier as _sklearn_ExtraTreesClassifier
+from sklearn.ensemble import ExtraTreesRegressor as _sklearn_ExtraTreesRegressor
+from sklearn.ensemble import RandomForestClassifier as _sklearn_RandomForestClassifier
+from sklearn.ensemble import RandomForestRegressor as _sklearn_RandomForestRegressor
+from sklearn.ensemble._forest import ForestClassifier as _sklearn_ForestClassifier
+from sklearn.ensemble._forest import ForestRegressor as _sklearn_ForestRegressor
 from sklearn.ensemble._forest import _get_n_samples_bootstrap
 from sklearn.exceptions import DataConversionWarning
+from sklearn.metrics import accuracy_score, r2_score
 from sklearn.tree import (
     DecisionTreeClassifier,
     DecisionTreeRegressor,
@@ -36,71 +39,59 @@ from sklearn.tree import (
 from sklearn.tree._tree import Tree
 from sklearn.utils import check_random_state, deprecated
 from sklearn.utils.validation import (
+    _check_sample_weight,
     check_array,
-    check_consistent_length,
     check_is_fitted,
     check_X_y,
 )
 
+from daal4py.sklearn._n_jobs_support import control_n_jobs
 from daal4py.sklearn._utils import (
     check_tree_nodes,
-    control_n_jobs,
     daal_check_version,
-    run_with_n_jobs,
     sklearn_check_version,
 )
 from onedal.ensemble import ExtraTreesClassifier as onedal_ExtraTreesClassifier
 from onedal.ensemble import ExtraTreesRegressor as onedal_ExtraTreesRegressor
 from onedal.ensemble import RandomForestClassifier as onedal_RandomForestClassifier
 from onedal.ensemble import RandomForestRegressor as onedal_RandomForestRegressor
-
-# try catch needed for changes in structures observed in Scikit-learn around v0.22
-try:
-    from sklearn.ensemble._forest import ForestClassifier as sklearn_ForestClassifier
-    from sklearn.ensemble._forest import ForestRegressor as sklearn_ForestRegressor
-except ModuleNotFoundError:
-    from sklearn.ensemble.forest import ForestClassifier as sklearn_ForestClassifier
-    from sklearn.ensemble.forest import ForestRegressor as sklearn_ForestRegressor
-
 from onedal.primitives import get_tree_state_cls, get_tree_state_reg
 from onedal.utils import _num_features, _num_samples
+from sklearnex import get_hyperparameters
+from sklearnex._utils import register_hyperparameters
 
-from .._config import get_config
 from .._device_offload import dispatch, wrap_output_data
 from .._utils import PatchingConditionsChain
+from ..utils._array_api import get_namespace
 
 if sklearn_check_version("1.2"):
     from sklearn.utils._param_validation import Interval
 if sklearn_check_version("1.4"):
     from daal4py.sklearn.utils import _assert_all_finite
 
+if sklearn_check_version("1.6"):
+    from sklearn.utils.validation import validate_data
+else:
+    validate_data = BaseEstimator._validate_data
+
 
 class BaseForest(ABC):
     _onedal_factory = None
 
-    @run_with_n_jobs
     def _onedal_fit(self, X, y, sample_weight=None, queue=None):
-        if sklearn_check_version("0.24"):
-            X, y = self._validate_data(
-                X,
-                y,
-                multi_output=False,
-                accept_sparse=False,
-                dtype=[np.float64, np.float32],
-                force_all_finite=False,
-            )
-        else:
-            X, y = check_X_y(
-                X,
-                y,
-                accept_sparse=False,
-                dtype=[np.float64, np.float32],
-                multi_output=False,
-                force_all_finite=False,
-            )
+        X, y = validate_data(
+            self,
+            X,
+            y,
+            multi_output=True,
+            accept_sparse=False,
+            dtype=[np.float64, np.float32],
+            force_all_finite=False,
+            ensure_2d=True,
+        )
 
         if sample_weight is not None:
-            sample_weight = self.check_sample_weight(sample_weight, X)
+            sample_weight = _check_sample_weight(sample_weight, X)
 
         if y.ndim == 2 and y.shape[1] == 1:
             warnings.warn(
@@ -120,8 +111,6 @@ class BaseForest(ABC):
 
         y, expanded_class_weight = self._validate_y_class_weight(y)
 
-        self.n_features_in_ = X.shape[1]
-
         if expanded_class_weight is not None:
             if sample_weight is not None:
                 sample_weight = sample_weight * expanded_class_weight
@@ -137,7 +126,9 @@ class BaseForest(ABC):
             "min_samples_split": self.min_samples_split,
             "min_samples_leaf": self.min_samples_leaf,
             "min_weight_fraction_leaf": self.min_weight_fraction_leaf,
-            "max_features": self.max_features,
+            "max_features": self._to_absolute_max_features(
+                self.max_features, self.n_features_in_
+            ),
             "max_leaf_nodes": self.max_leaf_nodes,
             "min_impurity_decrease": self.min_impurity_decrease,
             "bootstrap": self.bootstrap,
@@ -175,15 +166,6 @@ class BaseForest(ABC):
 
         return self
 
-    def _fit_proba(self, X, y, sample_weight=None, queue=None):
-        params = self.get_params()
-        self.__class__(**params)
-
-        # We use stock metaestimators below, so the only way
-        # to pass a queue is using config_context.
-        cfg = get_config()
-        cfg["target_offload"] = queue
-
     def _save_attributes(self):
         if self.oob_score:
             self.oob_score_ = self._onedal_estimator.oob_score_
@@ -206,8 +188,45 @@ class BaseForest(ABC):
         self._validate_estimator()
         return self
 
-    # TODO:
-    # move to onedal modul.
+    def _to_absolute_max_features(self, max_features, n_features):
+        if max_features is None:
+            return n_features
+        if isinstance(max_features, str):
+            if max_features == "auto":
+                if not sklearn_check_version("1.3"):
+                    if sklearn_check_version("1.1"):
+                        warnings.warn(
+                            "`max_features='auto'` has been deprecated in 1.1 "
+                            "and will be removed in 1.3. To keep the past behaviour, "
+                            "explicitly set `max_features=1.0` or remove this "
+                            "parameter as it is also the default value for "
+                            "RandomForestRegressors and ExtraTreesRegressors.",
+                            FutureWarning,
+                        )
+                    return (
+                        max(1, int(np.sqrt(n_features)))
+                        if isinstance(self, ForestClassifier)
+                        else n_features
+                    )
+            if max_features == "sqrt":
+                return max(1, int(np.sqrt(n_features)))
+            if max_features == "log2":
+                return max(1, int(np.log2(n_features)))
+            allowed_string_values = (
+                '"sqrt" or "log2"'
+                if sklearn_check_version("1.3")
+                else '"auto", "sqrt" or "log2"'
+            )
+            raise ValueError(
+                "Invalid value for max_features. Allowed string "
+                f"values are {allowed_string_values}."
+            )
+        if isinstance(max_features, (numbers.Integral, np.integer)):
+            return max_features
+        if max_features > 0.0:
+            return max(1, int(max_features * n_features))
+        return 0
+
     def _check_parameters(self):
         if isinstance(self.min_samples_leaf, numbers.Integral):
             if not 1 <= self.min_samples_leaf:
@@ -283,38 +302,6 @@ class BaseForest(ABC):
                 "min_bin_size must be integral number but was " "%r" % self.min_bin_size
             )
 
-    def check_sample_weight(self, sample_weight, X, dtype=None):
-        n_samples = _num_samples(X)
-
-        if dtype is not None and dtype not in [np.float32, np.float64]:
-            dtype = np.float64
-
-        if sample_weight is None:
-            sample_weight = np.ones(n_samples, dtype=dtype)
-        elif isinstance(sample_weight, numbers.Number):
-            sample_weight = np.full(n_samples, sample_weight, dtype=dtype)
-        else:
-            if dtype is None:
-                dtype = [np.float64, np.float32]
-            sample_weight = check_array(
-                sample_weight,
-                accept_sparse=False,
-                ensure_2d=False,
-                dtype=dtype,
-                order="C",
-                force_all_finite=False,
-            )
-            if sample_weight.ndim != 1:
-                raise ValueError("Sample weights must be 1D array or scalar")
-
-            if sample_weight.shape != (n_samples,):
-                raise ValueError(
-                    "sample_weight.shape == {}, expected {}!".format(
-                        sample_weight.shape, (n_samples,)
-                    )
-                )
-        return sample_weight
-
     @property
     def estimators_(self):
         if hasattr(self, "_cached_estimators_"):
@@ -415,7 +402,7 @@ class BaseForest(ABC):
             self.estimator = estimator
 
 
-class ForestClassifier(sklearn_ForestClassifier, BaseForest):
+class ForestClassifier(_sklearn_ForestClassifier, BaseForest):
     # Surprisingly, even though scikit-learn warns against using
     # their ForestClassifier directly, it actually has a more stable
     # API than the user-facing objects (over time). If they change it
@@ -455,14 +442,12 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
 
         # The estimator is checked against the class attribute for conformance.
         # This should only trigger if the user uses this class directly.
-        if (
-            self.estimator.__class__ == DecisionTreeClassifier
-            and self._onedal_factory != onedal_RandomForestClassifier
+        if self.estimator.__class__ == DecisionTreeClassifier and not issubclass(
+            self._onedal_factory, onedal_RandomForestClassifier
         ):
             self._onedal_factory = onedal_RandomForestClassifier
-        elif (
-            self.estimator.__class__ == ExtraTreeClassifier
-            and self._onedal_factory != onedal_ExtraTreesClassifier
+        elif self.estimator.__class__ == ExtraTreeClassifier and not issubclass(
+            self._onedal_factory, onedal_ExtraTreesClassifier
         ):
             self._onedal_factory = onedal_ExtraTreesClassifier
 
@@ -481,7 +466,7 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
             "fit",
             {
                 "onedal": self.__class__._onedal_fit,
-                "sklearn": sklearn_ForestClassifier.fit,
+                "sklearn": _sklearn_ForestClassifier.fit,
             },
             X,
             y,
@@ -554,18 +539,14 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
             )
 
         if patching_status.get_status():
-            if sklearn_check_version("0.24"):
-                X, y = self._validate_data(
-                    X,
-                    y,
-                    multi_output=True,
-                    accept_sparse=True,
-                    dtype=[np.float64, np.float32],
-                    force_all_finite=False,
-                )
-            else:
-                X = check_array(X, dtype=[np.float64, np.float32], force_all_finite=False)
-                y = check_array(y, ensure_2d=False, dtype=X.dtype, force_all_finite=False)
+            X, y = check_X_y(
+                X,
+                y,
+                multi_output=True,
+                accept_sparse=True,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+            )
 
             if y.ndim == 2 and y.shape[1] == 1:
                 warnings.warn(
@@ -619,12 +600,13 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
 
     @wrap_output_data
     def predict(self, X):
+        check_is_fitted(self)
         return dispatch(
             self,
             "predict",
             {
                 "onedal": self.__class__._onedal_predict,
-                "sklearn": sklearn_ForestClassifier.predict,
+                "sklearn": _sklearn_ForestClassifier.predict,
             },
             X,
         )
@@ -634,34 +616,50 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
         # TODO:
         # _check_proba()
         # self._check_proba()
-        if sklearn_check_version("1.0"):
-            self._check_feature_names(X, reset=False)
-        if hasattr(self, "n_features_in_"):
-            try:
-                num_features = _num_features(X)
-            except TypeError:
-                num_features = _num_samples(X)
-            if num_features != self.n_features_in_:
-                raise ValueError(
-                    (
-                        f"X has {num_features} features, "
-                        f"but {self.__class__.__name__} is expecting "
-                        f"{self.n_features_in_} features as input"
-                    )
-                )
+        check_is_fitted(self)
         return dispatch(
             self,
             "predict_proba",
             {
                 "onedal": self.__class__._onedal_predict_proba,
-                "sklearn": sklearn_ForestClassifier.predict_proba,
+                "sklearn": _sklearn_ForestClassifier.predict_proba,
             },
             X,
         )
 
-    fit.__doc__ = sklearn_ForestClassifier.fit.__doc__
-    predict.__doc__ = sklearn_ForestClassifier.predict.__doc__
-    predict_proba.__doc__ = sklearn_ForestClassifier.predict_proba.__doc__
+    def predict_log_proba(self, X):
+        xp, _ = get_namespace(X)
+        proba = self.predict_proba(X)
+
+        if self.n_outputs_ == 1:
+            return xp.log(proba)
+
+        else:
+            for k in range(self.n_outputs_):
+                proba[k] = xp.log(proba[k])
+
+            return proba
+
+    @wrap_output_data
+    def score(self, X, y, sample_weight=None):
+        check_is_fitted(self)
+        return dispatch(
+            self,
+            "score",
+            {
+                "onedal": self.__class__._onedal_score,
+                "sklearn": _sklearn_ForestClassifier.score,
+            },
+            X,
+            y,
+            sample_weight=sample_weight,
+        )
+
+    fit.__doc__ = _sklearn_ForestClassifier.fit.__doc__
+    predict.__doc__ = _sklearn_ForestClassifier.predict.__doc__
+    predict_proba.__doc__ = _sklearn_ForestClassifier.predict_proba.__doc__
+    predict_log_proba.__doc__ = _sklearn_ForestClassifier.predict_log_proba.__doc__
+    score.__doc__ = _sklearn_ForestClassifier.score.__doc__
 
     def _onedal_cpu_supported(self, method_name, *data):
         class_name = self.__class__.__name__
@@ -688,7 +686,7 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
                 ]
             )
 
-        elif method_name in ["predict", "predict_proba"]:
+        elif method_name in ["predict", "predict_proba", "score"]:
             X = data[0]
 
             patching_status.and_conditions(
@@ -749,11 +747,15 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
                         or self.estimator.__class__ == DecisionTreeClassifier,
                         "ExtraTrees only supported starting from oneDAL version 2023.1",
                     ),
-                    (sample_weight is not None, "sample_weight is not supported."),
+                    (
+                        not self.oob_score,
+                        "oob_scores using r2 or accuracy not implemented.",
+                    ),
+                    (sample_weight is None, "sample_weight is not supported."),
                 ]
             )
 
-        elif method_name in ["predict", "predict_proba"]:
+        elif method_name in ["predict", "predict_proba", "score"]:
             X = data[0]
 
             patching_status.and_conditions(
@@ -787,34 +789,69 @@ class ForestClassifier(sklearn_ForestClassifier, BaseForest):
 
         return patching_status
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
-        X = check_array(
-            X,
-            dtype=[np.float64, np.float32],
-            force_all_finite=False,
-        )  # Warning, order of dtype matters
-        check_is_fitted(self, "_onedal_estimator")
 
         if sklearn_check_version("1.0"):
-            self._check_feature_names(X, reset=False)
+            X = validate_data(
+                self,
+                X,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+                reset=False,
+                ensure_2d=True,
+            )
+        else:
+            X = check_array(
+                X,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+            )  # Warning, order of dtype matters
+            if hasattr(self, "n_features_in_"):
+                try:
+                    num_features = _num_features(X)
+                except TypeError:
+                    num_features = _num_samples(X)
+                if num_features != self.n_features_in_:
+                    raise ValueError(
+                        (
+                            f"X has {num_features} features, "
+                            f"but {self.__class__.__name__} is expecting "
+                            f"{self.n_features_in_} features as input"
+                        )
+                    )
+            self._check_n_features(X, reset=False)
 
         res = self._onedal_estimator.predict(X, queue=queue)
         return np.take(self.classes_, res.ravel().astype(np.int64, casting="unsafe"))
 
-    @run_with_n_jobs
     def _onedal_predict_proba(self, X, queue=None):
-        X = check_array(X, dtype=[np.float64, np.float32], force_all_finite=False)
-        check_is_fitted(self, "_onedal_estimator")
 
-        if sklearn_check_version("0.23"):
-            self._check_n_features(X, reset=False)
         if sklearn_check_version("1.0"):
-            self._check_feature_names(X, reset=False)
+            X = validate_data(
+                self,
+                X,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+                reset=False,
+                ensure_2d=True,
+            )
+        else:
+            X = check_array(
+                X,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+            )  # Warning, order of dtype matters
+            self._check_n_features(X, reset=False)
+
         return self._onedal_estimator.predict_proba(X, queue=queue)
 
+    def _onedal_score(self, X, y, sample_weight=None, queue=None):
+        return accuracy_score(
+            y, self._onedal_predict(X, queue=queue), sample_weight=sample_weight
+        )
+
 
-class ForestRegressor(sklearn_ForestRegressor, BaseForest):
+class ForestRegressor(_sklearn_ForestRegressor, BaseForest):
     _err = "out_of_bag_error_r2|out_of_bag_error_prediction"
     _get_tree_state = staticmethod(get_tree_state_reg)
 
@@ -847,14 +884,12 @@ class ForestRegressor(sklearn_ForestRegressor, BaseForest):
 
         # The splitter is checked against the class attribute for conformance
         # This should only trigger if the user uses this class directly.
-        if (
-            self.estimator.__class__ == DecisionTreeRegressor
-            and self._onedal_factory != onedal_RandomForestRegressor
+        if self.estimator.__class__ == DecisionTreeRegressor and not issubclass(
+            self._onedal_factory, onedal_RandomForestRegressor
         ):
             self._onedal_factory = onedal_RandomForestRegressor
-        elif (
-            self.estimator.__class__ == ExtraTreeRegressor
-            and self._onedal_factory != onedal_ExtraTreesRegressor
+        elif self.estimator.__class__ == ExtraTreeRegressor and not issubclass(
+            self._onedal_factory, onedal_ExtraTreesRegressor
         ):
             self._onedal_factory = onedal_ExtraTreesRegressor
 
@@ -924,18 +959,14 @@ class ForestRegressor(sklearn_ForestRegressor, BaseForest):
             )
 
         if patching_status.get_status():
-            if sklearn_check_version("0.24"):
-                X, y = self._validate_data(
-                    X,
-                    y,
-                    multi_output=True,
-                    accept_sparse=True,
-                    dtype=[np.float64, np.float32],
-                    force_all_finite=False,
-                )
-            else:
-                X = check_array(X, dtype=[np.float64, np.float32], force_all_finite=False)
-                y = check_array(y, ensure_2d=False, dtype=X.dtype, force_all_finite=False)
+            X, y = check_X_y(
+                X,
+                y,
+                multi_output=True,
+                accept_sparse=True,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+            )
 
             if y.ndim == 2 and y.shape[1] == 1:
                 warnings.warn(
@@ -1010,7 +1041,7 @@ class ForestRegressor(sklearn_ForestRegressor, BaseForest):
                 ]
             )
 
-        elif method_name == "predict":
+        elif method_name in ["predict", "score"]:
             X = data[0]
 
             patching_status.and_conditions(
@@ -1060,11 +1091,12 @@ class ForestRegressor(sklearn_ForestRegressor, BaseForest):
                         or self.estimator.__class__ == DecisionTreeClassifier,
                         "ExtraTrees only supported starting from oneDAL version 2023.1",
                     ),
-                    (sample_weight is not None, "sample_weight is not supported."),
+                    (not self.oob_score, "oob_score value is not sklearn conformant."),
+                    (sample_weight is None, "sample_weight is not supported."),
                 ]
             )
 
-        elif method_name == "predict":
+        elif method_name in ["predict", "score"]:
             X = data[0]
 
             patching_status.and_conditions(
@@ -1096,25 +1128,37 @@ class ForestRegressor(sklearn_ForestRegressor, BaseForest):
 
         return patching_status
 
-    @run_with_n_jobs
     def _onedal_predict(self, X, queue=None):
-        X = check_array(
-            X, dtype=[np.float64, np.float32], force_all_finite=False
-        )  # Warning, order of dtype matters
         check_is_fitted(self, "_onedal_estimator")
 
         if sklearn_check_version("1.0"):
-            self._check_feature_names(X, reset=False)
+            X = validate_data(
+                self,
+                X,
+                dtype=[np.float64, np.float32],
+                force_all_finite=False,
+                reset=False,
+                ensure_2d=True,
+            )  # Warning, order of dtype matters
+        else:
+            X = check_array(
+                X, dtype=[np.float64, np.float32], force_all_finite=False
+            )  # Warning, order of dtype matters
 
         return self._onedal_estimator.predict(X, queue=queue)
 
+    def _onedal_score(self, X, y, sample_weight=None, queue=None):
+        return r2_score(
+            y, self._onedal_predict(X, queue=queue), sample_weight=sample_weight
+        )
+
     def fit(self, X, y, sample_weight=None):
         dispatch(
             self,
             "fit",
             {
                 "onedal": self.__class__._onedal_fit,
-                "sklearn": sklearn_ForestRegressor.fit,
+                "sklearn": _sklearn_ForestRegressor.fit,
             },
             X,
             y,
@@ -1124,28 +1168,46 @@ class ForestRegressor(sklearn_ForestRegressor, BaseForest):
 
     @wrap_output_data
     def predict(self, X):
+        check_is_fitted(self)
         return dispatch(
             self,
             "predict",
             {
                 "onedal": self.__class__._onedal_predict,
-                "sklearn": sklearn_ForestRegressor.predict,
+                "sklearn": _sklearn_ForestRegressor.predict,
             },
             X,
         )
 
-    fit.__doc__ = sklearn_ForestRegressor.fit.__doc__
-    predict.__doc__ = sklearn_ForestRegressor.predict.__doc__
+    @wrap_output_data
+    def score(self, X, y, sample_weight=None):
+        check_is_fitted(self)
+        return dispatch(
+            self,
+            "score",
+            {
+                "onedal": self.__class__._onedal_score,
+                "sklearn": _sklearn_ForestRegressor.score,
+            },
+            X,
+            y,
+            sample_weight=sample_weight,
+        )
+
+    fit.__doc__ = _sklearn_ForestRegressor.fit.__doc__
+    predict.__doc__ = _sklearn_ForestRegressor.predict.__doc__
+    score.__doc__ = _sklearn_ForestRegressor.score.__doc__
 
 
-@control_n_jobs
+@register_hyperparameters({"infer": get_hyperparameters("decision_forest", "infer")})
+@control_n_jobs(decorated_methods=["fit", "predict", "predict_proba", "score"])
 class RandomForestClassifier(ForestClassifier):
-    __doc__ = sklearn_RandomForestClassifier.__doc__
+    __doc__ = _sklearn_RandomForestClassifier.__doc__
     _onedal_factory = onedal_RandomForestClassifier
 
     if sklearn_check_version("1.2"):
         _parameter_constraints: dict = {
-            **sklearn_RandomForestClassifier._parameter_constraints,
+            **_sklearn_RandomForestClassifier._parameter_constraints,
             "max_bins": [Interval(numbers.Integral, 2, None, closed="left")],
             "min_bin_size": [Interval(numbers.Integral, 1, None, closed="left")],
         }
@@ -1348,14 +1410,14 @@ class RandomForestClassifier(ForestClassifier):
             self.min_bin_size = min_bin_size
 
 
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "predict", "score"])
 class RandomForestRegressor(ForestRegressor):
-    __doc__ = sklearn_RandomForestRegressor.__doc__
+    __doc__ = _sklearn_RandomForestRegressor.__doc__
     _onedal_factory = onedal_RandomForestRegressor
 
     if sklearn_check_version("1.2"):
         _parameter_constraints: dict = {
-            **sklearn_RandomForestRegressor._parameter_constraints,
+            **_sklearn_RandomForestRegressor._parameter_constraints,
             "max_bins": [Interval(numbers.Integral, 2, None, closed="left")],
             "min_bin_size": [Interval(numbers.Integral, 1, None, closed="left")],
         }
@@ -1549,14 +1611,14 @@ class RandomForestRegressor(ForestRegressor):
             self.min_bin_size = min_bin_size
 
 
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "predict", "predict_proba", "score"])
 class ExtraTreesClassifier(ForestClassifier):
-    __doc__ = sklearn_ExtraTreesClassifier.__doc__
+    __doc__ = _sklearn_ExtraTreesClassifier.__doc__
     _onedal_factory = onedal_ExtraTreesClassifier
 
     if sklearn_check_version("1.2"):
         _parameter_constraints: dict = {
-            **sklearn_ExtraTreesClassifier._parameter_constraints,
+            **_sklearn_ExtraTreesClassifier._parameter_constraints,
             "max_bins": [Interval(numbers.Integral, 2, None, closed="left")],
             "min_bin_size": [Interval(numbers.Integral, 1, None, closed="left")],
         }
@@ -1759,14 +1821,14 @@ class ExtraTreesClassifier(ForestClassifier):
             self.min_bin_size = min_bin_size
 
 
-@control_n_jobs
+@control_n_jobs(decorated_methods=["fit", "predict", "score"])
 class ExtraTreesRegressor(ForestRegressor):
-    __doc__ = sklearn_ExtraTreesRegressor.__doc__
+    __doc__ = _sklearn_ExtraTreesRegressor.__doc__
     _onedal_factory = onedal_ExtraTreesRegressor
 
     if sklearn_check_version("1.2"):
         _parameter_constraints: dict = {
-            **sklearn_ExtraTreesRegressor._parameter_constraints,
+            **_sklearn_ExtraTreesRegressor._parameter_constraints,
             "max_bins": [Interval(numbers.Integral, 2, None, closed="left")],
             "min_bin_size": [Interval(numbers.Integral, 1, None, closed="left")],
         }
@@ -1961,7 +2023,7 @@ class ExtraTreesRegressor(ForestRegressor):
 
 
 # Allow for isinstance calls without inheritance changes using ABCMeta
-sklearn_RandomForestClassifier.register(RandomForestClassifier)
-sklearn_RandomForestRegressor.register(RandomForestRegressor)
-sklearn_ExtraTreesClassifier.register(ExtraTreesClassifier)
-sklearn_ExtraTreesRegressor.register(ExtraTreesRegressor)
+_sklearn_RandomForestClassifier.register(RandomForestClassifier)
+_sklearn_RandomForestRegressor.register(RandomForestRegressor)
+_sklearn_ExtraTreesClassifier.register(ExtraTreesClassifier)
+_sklearn_ExtraTreesRegressor.register(ExtraTreesRegressor)