Fix multi-output

Signed-off-by: Adam Li <adam2392@gmail.com>
neurodata · Jun 15, 2023 · f1ab592 · f1ab592
1 parent f831bc4
commit f1ab592
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 7 deletions.
diff --git a/sktree/_lib/sklearn_fork b/sktree/_lib/sklearn_fork
diff --git a/sktree/ensemble/_honest_forest.py b/sktree/ensemble/_honest_forest.py
@@ -451,6 +451,7 @@ def _predict_proba(self, X, indices=None, impute_missing=None):
             delayed(_accumulate_prediction)(tree, X, posteriors, lock, idx)
             for tree, idx in zip(self.estimators_, indices)
         )
+
         # Normalize to unit length, due to prior weighting
         posteriors = np.array(posteriors)
         zero_mask = posteriors.sum(2) == 0
@@ -461,6 +462,10 @@ def _predict_proba(self, X, indices=None, impute_missing=None):
         else:
             posteriors[zero_mask] = impute_missing
 
+        # preserve shape of multi-outputs
+        if self.n_outputs_ > 1:
+            posteriors = [post for post in posteriors]
+
         if len(posteriors) == 1:
             return posteriors[0]
         else:

diff --git a/sktree/tests/test_honest_forest.py b/sktree/tests/test_honest_forest.py
@@ -138,14 +138,12 @@ def test_honest_decision_function(honest_fraction, val):
     [HonestForestClassifier(n_estimators=10, honest_fraction=0.5, random_state=0)]
 )
 def test_sklearn_compatible_estimator(estimator, check):
-    # 1. multi-output is not fully supported
-    # 2. check_class_weight_classifiers is not supported since it requires sample weight
+    # 1. check_class_weight_classifiers is not supported since it requires sample weight
     # XXX: can include this "generalization" in the future if it's useful
     #  zero sample weight is not "really supported" in honest subsample trees since sample weight
     #  for fitting the tree's splits
     if check.func.__name__ in [
         "check_class_weight_classifiers",
-        "check_classifiers_multilabel_output_format_predict_proba",
     ]:
         pytest.skip()
     check(estimator)
diff --git a/sktree/tree/tests/test_honest_tree.py b/sktree/tree/tests/test_honest_tree.py
@@ -106,7 +106,4 @@ def test_impute_classes():
 
 @parametrize_with_checks([HonestTreeClassifier(random_state=0)])
 def test_sklearn_compatible_estimator(estimator, check):
-    # TODO: remove when we implement Regressor classes
-    # if TREE_ESTIMATORS[estimator].__name__ in TREE_CLASSIFIERS:
-    #     pytest.skip()
     check(estimator)