From 721173c291396c28b355b93f2c3eb6a571beaa19 Mon Sep 17 00:00:00 2001 From: Jeremy Schoemaker Date: Mon, 24 Aug 2026 03:02:43 -0500 Subject: [PATCH 1/2] fix: split sample_weight in group holdout instead of passing the full vector With split_type="group", eval_method="holdout", and sample_weight set, the group branch of _prepare_data splits X, y, and groups by train_idx/val_idx but leaves state.fit_kwargs["sample_weight"] at full length and never sets state.weight_val. AutoMLState.prepare_sample_train_data then slices the weights positionally (weight[:sample_size]), so every trial trains with weights belonging to other rows, and the validation loss that drives model selection ignores sample weights entirely. No error is raised. The "time" branch (generic_task.py:913-955) and _train_test_split (generic_task.py:309) already split the weights this way; the group branch is the only holdout branch that does not. This mirrors that pattern inside the gss loop: slice the weights by train_idx/val_idx into state.fit_kwargs["sample_weight"] and state.weight_val, handling both pd.Series and array inputs. The regression test spies on RandomForestClassifier.fit and encodes each row's expected weight in the row itself, so any misalignment between the rows an estimator receives and the weights it receives fails the assert. --- flaml/automl/task/generic_task.py | 9 ++++++ test/automl/test_split.py | 51 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/flaml/automl/task/generic_task.py b/flaml/automl/task/generic_task.py index 0726332ba2..bb2afd71c1 100644 --- a/flaml/automl/task/generic_task.py +++ b/flaml/automl/task/generic_task.py @@ -964,6 +964,15 @@ def prepare_data( y_train, y_val = y_train_all[train_idx], y_train_all[val_idx] state.groups = state.groups_all[train_idx] state.groups_val = state.groups_all[val_idx] + if "sample_weight" in state.fit_kwargs: + # NOTE: _prepare_data is before kwargs is updated to fit_kwargs_by_estimator + weight = state.fit_kwargs["sample_weight"] + if isinstance(weight, pd.Series): + state.fit_kwargs["sample_weight"] = weight.iloc[train_idx] + state.weight_val = weight.iloc[val_idx] + else: + state.fit_kwargs["sample_weight"] = weight[train_idx] + state.weight_val = weight[val_idx] elif self.is_classification(): # for classification, make sure the labels are complete in both # training and validation data diff --git a/test/automl/test_split.py b/test/automl/test_split.py index 8eba4bf28d..b3dab66e7a 100644 --- a/test/automl/test_split.py +++ b/test/automl/test_split.py @@ -74,6 +74,57 @@ def test_time_split_with_sample_weight(): assert automl.model is not None +def test_group_split_with_sample_weight_alignment(): + """Regression for group holdout: sample_weight must be split by the group + train/val indices, not left full-length and sliced positionally downstream.""" + from unittest.mock import patch + + from sklearn.ensemble import RandomForestClassifier + + n = 200 + rng = np.random.default_rng(42) + # column 0 is a row id, so each row's expected weight is recoverable from X + X = np.column_stack([np.arange(n, dtype=float), rng.normal(size=n)]) + groups = np.repeat(np.arange(20), 10) + y = (groups % 2).astype(int) + sample_weight = 1000.0 + np.arange(n, dtype=float) + + captured = [] + original_fit = RandomForestClassifier.fit + + def spy_fit(self, X, y, sample_weight=None, **kwargs): + captured.append( + ( + np.asarray(X).copy(), + None if sample_weight is None else np.asarray(sample_weight, dtype=float).copy(), + ) + ) + return original_fit(self, X, y, sample_weight=sample_weight, **kwargs) + + automl = AutoML() + with patch.object(RandomForestClassifier, "fit", spy_fit): + automl.fit( + X_train=X, + y_train=y, + sample_weight=sample_weight, + groups=groups, + split_type="group", + eval_method="holdout", + task="classification", + time_budget=-1, + max_iter=2, + estimator_list=["rf"], + ) + assert automl.model is not None + assert captured + for X_fit, weight_fit in captured: + assert weight_fit is not None + # each row trains with its own weight (1000 + row id) after the split + np.testing.assert_array_equal(weight_fit, 1000.0 + np.asarray(X_fit)[:, 0]) + # the validation loss must be weighted as well + assert automl._state.weight_val is not None + + def test_groups_for_classification_task(): from sklearn.externals._arff import ArffException From ddd01f13a037d0e62a56694d26b941fd586462ac Mon Sep 17 00:00:00 2001 From: Jeremy Schoemaker Date: Thu, 27 Aug 2026 15:45:35 -0500 Subject: [PATCH 2/2] fix: guard array sample_weight indexing and add Series regression test Convert non-Series sample_weight to ndarray before indexing by train_idx/ val_idx in the group holdout branch, so a plain list or tuple weight does not raise TypeError on fancy indexing, and skip splitting when weight is None. Add a regression test that uses a pd.Series with a non-default index to confirm the split stays positional (.iloc), mirroring the existing ndarray-only test. --- flaml/automl/task/generic_task.py | 7 +++- test/automl/test_split.py | 55 +++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/flaml/automl/task/generic_task.py b/flaml/automl/task/generic_task.py index bb2afd71c1..3c7cc7a1ed 100644 --- a/flaml/automl/task/generic_task.py +++ b/flaml/automl/task/generic_task.py @@ -967,10 +967,15 @@ def prepare_data( if "sample_weight" in state.fit_kwargs: # NOTE: _prepare_data is before kwargs is updated to fit_kwargs_by_estimator weight = state.fit_kwargs["sample_weight"] - if isinstance(weight, pd.Series): + if weight is None: + pass + elif isinstance(weight, pd.Series): state.fit_kwargs["sample_weight"] = weight.iloc[train_idx] state.weight_val = weight.iloc[val_idx] else: + # train_idx/val_idx are ndarrays from GroupShuffleSplit, so a plain + # list or tuple weight needs converting before it supports fancy indexing + weight = np.asarray(weight) state.fit_kwargs["sample_weight"] = weight[train_idx] state.weight_val = weight[val_idx] elif self.is_classification(): diff --git a/test/automl/test_split.py b/test/automl/test_split.py index b3dab66e7a..36d79fee87 100644 --- a/test/automl/test_split.py +++ b/test/automl/test_split.py @@ -125,6 +125,61 @@ def spy_fit(self, X, y, sample_weight=None, **kwargs): assert automl._state.weight_val is not None +def test_group_split_with_sample_weight_series_alignment(): + """Same regression as above, but sample_weight is a pd.Series with a + non-default index, to verify the group branch slices it positionally + (.iloc) rather than by label.""" + from unittest.mock import patch + + from sklearn.ensemble import RandomForestClassifier + + n = 200 + rng = np.random.default_rng(7) + X = pd.DataFrame( + {"id": np.arange(n, dtype=float), "feat": rng.normal(size=n)}, + ) + groups = np.repeat(np.arange(20), 10) + y = pd.Series((groups % 2).astype(int)) + # non-default index: label != positional order, so .loc and .iloc would disagree + weight_index = np.arange(n) + 500 + sample_weight = pd.Series(1000.0 + np.arange(n, dtype=float), index=weight_index) + + captured = [] + original_fit = RandomForestClassifier.fit + + def spy_fit(self, X, y, sample_weight=None, **kwargs): + captured.append( + ( + np.asarray(X).copy(), + None if sample_weight is None else np.asarray(sample_weight, dtype=float).copy(), + ) + ) + return original_fit(self, X, y, sample_weight=sample_weight, **kwargs) + + automl = AutoML() + with patch.object(RandomForestClassifier, "fit", spy_fit): + automl.fit( + X_train=X, + y_train=y, + sample_weight=sample_weight, + groups=groups, + split_type="group", + eval_method="holdout", + task="classification", + time_budget=-1, + max_iter=2, + estimator_list=["rf"], + ) + assert automl.model is not None + assert captured + for X_fit, weight_fit in captured: + assert weight_fit is not None + # each row trains with its own weight (1000 + row id) after the split, + # regardless of the Series' original (non-default) index + np.testing.assert_array_equal(weight_fit, 1000.0 + np.asarray(X_fit)[:, 0]) + assert automl._state.weight_val is not None + + def test_groups_for_classification_task(): from sklearn.externals._arff import ArffException