Skip to content

Commit

Permalink
enh: re-enabling spmd rf interfaces (#1700) (#1711)
Browse files Browse the repository at this point in the history
* enh: re-enabling spmd rf interfaces

* restoring onedal_factor usage but adding __class__

* lint

* another attempt at it

* isinstance to issubclass and re-adding self

(cherry picked from commit bfa470b)

Co-authored-by: ethanglaser <42726565+ethanglaser@users.noreply.github.com>
  • Loading branch information
mergify[bot] and ethanglaser authored Feb 13, 2024
1 parent 6c72039 commit 54f64c2
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 22 deletions.
20 changes: 8 additions & 12 deletions sklearnex/ensemble/_forest.py
Original file line number Diff line number Diff line change
Expand Up @@ -453,14 +453,12 @@ def __init__(

# 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

Expand Down Expand Up @@ -843,14 +841,12 @@ def __init__(

# 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

Expand Down
14 changes: 4 additions & 10 deletions sklearnex/spmd/ensemble/forest.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,9 @@
from ...ensemble import RandomForestRegressor as RandomForestRegressor_Batch


class BaseForestSPMD(ABC):
def _onedal_classifier(self, **onedal_params):
return onedal_RandomForestClassifier(**onedal_params)

def _onedal_regressor(self, **onedal_params):
return onedal_RandomForestRegressor(**onedal_params)


class RandomForestClassifier(BaseForestSPMD, RandomForestClassifier_Batch):
class RandomForestClassifier(RandomForestClassifier_Batch):
__doc__ = RandomForestClassifier_Batch.__doc__
_onedal_factory = onedal_RandomForestClassifier

def _onedal_cpu_supported(self, method_name, *data):
# TODO:
Expand All @@ -55,8 +48,9 @@ def _onedal_gpu_supported(self, method_name, *data):
return ready


class RandomForestRegressor(BaseForestSPMD, RandomForestRegressor_Batch):
class RandomForestRegressor(RandomForestRegressor_Batch):
__doc__ = RandomForestRegressor_Batch.__doc__
_onedal_factory = onedal_RandomForestRegressor

def _onedal_cpu_supported(self, method_name, *data):
# TODO:
Expand Down

0 comments on commit 54f64c2

Please sign in to comment.