pass extra kwargs from search.score to search.best_learner_.score

Author: jeromedockes

skrub/_data_ops/_estimator.py

@@ -1268,7 +1268,7 @@ def f(*args, **kwargs):
         f.__name__ = name
         return f
 
-    def _call_predictor_method(self, name, environment):
+    def _call_predictor_method(self, name, environment, **kwargs):
         check_is_fitted(self, "cv_results_")
         if not hasattr(self, "best_learner_"):
             raise AttributeError(
@@ -1277,7 +1277,7 @@ def _call_predictor_method(self, name, environment):
                 "Please pass another value to `refit` or fit a learner manually "
                 "using the `best_params_` or `cv_results_` attributes."
             )
-        return getattr(self.best_learner_, name)(environment)
+        return getattr(self.best_learner_, name)(environment, **kwargs)
 
     @property
     def results_(self):

skrub/_data_ops/tests/test_estimators.py

@@ -585,6 +585,10 @@ def test_randomized_search_with_scoring(randomized_search_backend):
     search.fit(split["train"])
     assert (search.results_["mean_test_neg_brier_score"] <= 0).all()
     assert list(search.score(split["test"]).keys()) == ["neg_brier_score", "roc_auc"]
+    # Test return_predictions is forwarded to best_learner_.score
+    scores, predictions = search.score(split["test"], return_predictions=True)
+    assert list(scores.keys()) == ["neg_brier_score", "roc_auc"]
+    assert predictions["predict_proba"].shape == (split["y_test"].shape[0], 2)
 
 
 def test_grid_search(data_op, data, n_jobs):