From db5acfb6c3ef2f62990faee451000122d9d4e8d3 Mon Sep 17 00:00:00 2001 From: thomasATbayer Date: Wed, 2 Sep 2026 12:18:09 +0000 Subject: [PATCH 1/3] fix(optimization): support hold-out CV in Optuna tuning (#58) Skip TerminatorCallback and cross-validation score reporting when hold-out validation produces fewer than two fold scores. Retain early termination for multi-fold cross-validation and add regression coverage for both callback paths. --- src/mother/optimization/core.py | 31 ++++++++++----- test/unit/test_model_tuner.py | 69 ++++++++++++++++++++++++++++++++- 2 files changed, 90 insertions(+), 10 deletions(-) diff --git a/src/mother/optimization/core.py b/src/mother/optimization/core.py index befef5e..5d79851 100644 --- a/src/mother/optimization/core.py +++ b/src/mother/optimization/core.py @@ -13,7 +13,11 @@ import sklearn.metrics as skl_metrics import sklearn.model_selection as skl_model_sel from optuna.study import Study, StudyDirection -from optuna.terminator import TerminatorCallback, report_cross_validation_scores +from optuna.terminator import ( + Terminator, + TerminatorCallback, + report_cross_validation_scores, +) from sklearn.pipeline import Pipeline from mother import utils as mother_utils @@ -133,7 +137,7 @@ def __init__( tuning_direction: typing.Union[StudyDirection, str] = StudyDirection.MAXIMIZE, n_trials_optuna: int = 100, n_threads_optuna: int = 1, - n_startup_trials: int = 12, + n_startup_trials: int = 20, seed: int = 42, **kwargs, ): @@ -157,17 +161,20 @@ def __init__( self.study: typing.Optional[Study] = None - def get_callbacks(self): + def get_callbacks(self, cross_validation: typing.Optional[skl_model_sel.BaseCrossValidator] = None): """ Prepares and returns a list of callbacks for early stopping in Optuna optimization. If early stopping with Optuna is enabled and PyTorch is available, this method - will return a list containing a TerminatorCallback instance. If PyTorch is not - available, it will log a warning and return None. + will return a list containing a TerminatorCallback instance. + + It returns None in either of the following cases: + - PyTorch is not available (warning is logged) + - hold-out cross-validation is detected (fewer than 2 splits) Returns: typing.Optional[typing.List[TerminatorCallback]]: A list of TerminatorCallback - instances if early stopping is enabled and PyTorch is available, otherwise None. + instances when early stopping can be used, otherwise None. """ callbacks: typing.Optional[typing.List[TerminatorCallback]] = None if self.early_stopping_optuna: @@ -179,7 +186,12 @@ def get_callbacks(self): pip install mother[torch] or uv add mother[torch]""" ) else: - callbacks = [TerminatorCallback()] + if cross_validation is not None and cross_validation.get_n_splits() < 2: + module_logger.warning( + "Optuna early termination requires at least 2 CV splits; disabling callback for hold-out setup" + ) + return None + callbacks = [TerminatorCallback(terminator=Terminator(min_n_trials=40))] return callbacks @handle_metadata_routing @@ -254,7 +266,8 @@ def objective(trial: optuna.trial.Trial) -> float: gc.collect() module_logger.info(f"Trial {trial.number}, cv score: {cv_score}") cv_score_not_na: np.ndarray = cv_score[~np.isnan(cv_score)] - report_cross_validation_scores(trial, list(cv_score_not_na)) + if len(cv_score_not_na) > 1: + report_cross_validation_scores(trial, list(cv_score_not_na)) mean_cv_score: float = cv_score_not_na.mean() return mean_cv_score @@ -284,7 +297,7 @@ def objective(trial: optuna.trial.Trial) -> float: objective, n_trials=self.n_trials_optuna, gc_after_trial=True, - callbacks=self.get_callbacks(), + callbacks=self.get_callbacks(cross_validation=cross_validation), ) if default_parameters != {}: diff --git a/test/unit/test_model_tuner.py b/test/unit/test_model_tuner.py index bb71e29..5988022 100644 --- a/test/unit/test_model_tuner.py +++ b/test/unit/test_model_tuner.py @@ -7,7 +7,7 @@ from sklearn.datasets import make_blobs, make_classification from sklearn.linear_model import Lasso from sklearn.metrics import accuracy_score, make_scorer, mean_squared_error -from sklearn.model_selection import GroupKFold +from sklearn.model_selection import GroupKFold, KFold, PredefinedSplit from sklearn.pipeline import Pipeline from sklearn.preprocessing import MinMaxScaler @@ -18,6 +18,13 @@ ) from mother.optimization.core import MotherTuner +try: + import torch # noqa: F401 + + _TORCH_AVAILABLE = True +except (ImportError, OSError): + _TORCH_AVAILABLE = False + # define a simple model to run the tests with @pytest.fixture() @@ -267,3 +274,63 @@ def test_tune_mother_feature_selection_model_pipeline( evaluated = tuner.study.trials[0].params enqueued = model_pipeline.default_parameters() all(evaluated[k] == enqueued[k] for k in evaluated.keys() & enqueued.keys()) + + +@pytest.mark.skipif( + not _TORCH_AVAILABLE, + reason="Optuna early-stopping callback tests require the optional torch extra (mother[torch]).", +) +class TestGetCallbacks: + def test_holdout_cv_disables_early_termination_callback(self, caplog): + caplog.set_level("WARNING") + tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True) + holdout_cv = PredefinedSplit(test_fold=[-1, -1, 0, 0]) + + callbacks = tuner.get_callbacks(cross_validation=holdout_cv) + + assert callbacks is None + assert "requires at least 2 CV splits" in caplog.text + + def test_multi_split_cv_keeps_early_termination_callback(self): + tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True) + kfold_cv = KFold(n_splits=3, shuffle=True, random_state=42) + + callbacks = tuner.get_callbacks(cross_validation=kfold_cv) + + assert callbacks is not None + assert len(callbacks) == 1 + + def test_no_cv_argument_keeps_existing_behavior(self): + tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True) + + callbacks = tuner.get_callbacks() + + assert callbacks is not None + assert len(callbacks) == 1 + + def test_optimize_works_with_holdout_cv(self, lasso_pipeline, synthetic_data, caplog): + caplog.set_level("WARNING") + X, y, _ = synthetic_data + test_fold = np.full(len(X), -1) + test_fold[-20:] = 0 + holdout_cv = PredefinedSplit(test_fold=test_fold) + + tuner = MotherTuner( + scorer="neg_mean_squared_error", + early_stopping_optuna=True, + n_trials_optuna=1, + ) + + model = tuner.optimize( + lasso_pipeline, + X=X, + y=y, + hyperparameter_space_function=lasso_pipeline.get_hyperparameter_space, + default_parameters=lasso_pipeline.default_parameters(), + cross_validation=holdout_cv, + ) + + assert model is not None + assert tuner.study is not None + assert len(tuner.study.trials) >= 1 + assert "requires at least 2 CV splits" in caplog.text From 44fd13cc0acd327597fe539c88f55e3d174df462 Mon Sep 17 00:00:00 2001 From: thomasATbayer Date: Wed, 2 Sep 2026 14:02:52 +0000 Subject: [PATCH 2/3] fix(tabpfn): disable autocast for pre-fitted embeddings Force pre-fitted TabPFN embedding extraction to run without autocast before calling get_embeddings(). With Python 3.14 and the current TabPFN runtime, pre-fitted models can emit bfloat16 embedding tensors, which NumPy cannot convert and raises: TypeError: Got unsupported ScalarType BFloat16 This only affects the pre-fitted model path in TabPFNEmbeddingTransformer. Models fitted inside the transformer already pass inference_precision=torch.float32 during construction, so their embedding path is unchanged. Add a regression assertion to ensure the pre-fitted path disables autocast before embedding extraction. --- src/mother/ml/models/m_tabpfn.py | 8 ++++++++ test/unit/test_tabpfn.py | 1 + 2 files changed, 9 insertions(+) diff --git a/src/mother/ml/models/m_tabpfn.py b/src/mother/ml/models/m_tabpfn.py index eb872d7..3fc24f7 100644 --- a/src/mother/ml/models/m_tabpfn.py +++ b/src/mother/ml/models/m_tabpfn.py @@ -536,6 +536,11 @@ def _set_random_state(self) -> None: torch.manual_seed(self.random_state) np.random.seed(self.random_state) + def _prepare_prefitted_model_for_embeddings(self) -> None: + if self.model is None: + raise RuntimeError("A pre-fitted model is required to extract embeddings.") + self.model.use_autocast_ = False + def _get_best_embeddings( self, embeddings: np.ndarray, @@ -633,6 +638,7 @@ def fit( module_logger.info( "A pre-fitted model has been given. The new data will not be used for fitting the model." ) + self._prepare_prefitted_model_for_embeddings() self.train_embeddings_ = self.model.get_embeddings(X_array) self._embedding_dim = self.train_embeddings_.shape[1] else: @@ -777,6 +783,8 @@ def transform( X_array = np.asarray(X, dtype=np.float32) # Get embeddings for new data using the main model + if self.pre_fitted: + self._prepare_prefitted_model_for_embeddings() embeddings = self.model.get_embeddings(X_array) # collapse the additional column caused by estimators (avg) if len(embeddings.shape) == 3: diff --git a/test/unit/test_tabpfn.py b/test/unit/test_tabpfn.py index d2a5a14..29df275 100644 --- a/test/unit/test_tabpfn.py +++ b/test/unit/test_tabpfn.py @@ -208,6 +208,7 @@ def test_prefitted_model(self): ) result_prefitted = transformer_prefitted.transform(self.X) + assert transformer_prefitted.model.use_autocast_ is False transformer_newfit = TabPFNEmbeddingTransformer( model_type="regression", use_kfold=False, n_estimators=1, random_state=0 From 40ed38fcdf29dbd05eb38a72d214d7fdecfd55ba Mon Sep 17 00:00:00 2001 From: thomasATbayer Date: Wed, 2 Sep 2026 14:18:50 +0000 Subject: [PATCH 3/3] docs(optimization): clarify Optuna callback disablement cases Update the MotherTuner.get_callbacks docstring so its documented None return conditions match the implementation. The method can skip creating Optuna early-stopping callbacks when: - early_stopping_optuna is disabled - PyTorch is unavailable - hold-out cross-validation is used This prevents callers from assuming None only indicates a missing PyTorch dependency or hold-out CV path. --- src/mother/optimization/core.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/mother/optimization/core.py b/src/mother/optimization/core.py index 5d79851..8aaefbb 100644 --- a/src/mother/optimization/core.py +++ b/src/mother/optimization/core.py @@ -165,10 +165,12 @@ def get_callbacks(self, cross_validation: typing.Optional[skl_model_sel.BaseCros """ Prepares and returns a list of callbacks for early stopping in Optuna optimization. - If early stopping with Optuna is enabled and PyTorch is available, this method - will return a list containing a TerminatorCallback instance. + If early stopping with Optuna is enabled, PyTorch is available, and the + cross-validation strategy supports early termination, this method will return + a list containing a TerminatorCallback instance. - It returns None in either of the following cases: + It returns None in any of the following cases: + - early stopping with Optuna is disabled - PyTorch is not available (warning is logged) - hold-out cross-validation is detected (fewer than 2 splits)