From 41d6f93ed11a1708a6a589c9decaa9ca19331aec Mon Sep 17 00:00:00 2001 From: tamannarayan Date: Tue, 28 Jul 2026 12:07:08 -0400 Subject: [PATCH 1/4] Adopt shared hold-out set to optimize NNLS and calibration parameters on when per-member row subsampling is active. Allows the joint setting of max_num_rows along with NNLS and/or calibration. --- tabfm/src/classifier_and_regressor.py | 346 +++++++++++++++------ tabfm/src/classifier_and_regressor_test.py | 104 ++++++- 2 files changed, 351 insertions(+), 99 deletions(-) diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index eca99b1..e263d84 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/tabfm/src/classifier_and_regressor.py @@ -1063,6 +1063,8 @@ class EnsembleGenerator(TransformerMixin, BaseEstimator): svd_pool_end_: Ending column index in ``X_`` for the SVD features pool. k_crosses_list_: List of feature cross counts to sample per ensemble member. k_svd_list_: List of SVD feature counts to sample per ensemble member. + holdout_indices_: Array of absolute row indices in ``X_`` reserved for + global holdout validation when max_num_rows < dataset size. """ norm_methods_: List[str] @@ -1089,6 +1091,7 @@ class EnsembleGenerator(TransformerMixin, BaseEstimator): svd_pool_end_: int k_crosses_list_: List[int] k_svd_list_: List[int] + holdout_indices_: Optional[np.ndarray] def __init__( self, @@ -1106,6 +1109,7 @@ def __init__( total_svd_pool: Optional[int] = None, random_state: Optional[int] = None, task: str = "classification", + reserve_holdout_set: bool = False, ): """Initialises the generator. @@ -1130,6 +1134,8 @@ def __init__( ensemble member, or ``0`` to disable. total_svd_pool: Total pool size of SVD features to generate. task: Either ``"classification"`` or ``"regression"``. + reserve_holdout_set: Whether to reserve holdout validation rows when + max_num_rows < dataset size. """ self.n_estimators = n_estimators self.norm_methods = norm_methods @@ -1145,6 +1151,27 @@ def __init__( self.total_svd_pool = total_svd_pool self.random_state = random_state self.task = task + self.reserve_holdout_set = reserve_holdout_set + + @property + def holdout_indices(self) -> Optional[np.ndarray]: + """Public read-only accessor for global holdout validation row indices.""" + return self.holdout_indices_ + + def _compute_holdout_val_size(self, n_samples: int) -> int: + """Compute target holdout validation set size for row-subsampling regime. + + When max_num_rows is active alongside shared validation calculation, + we reserve a 5% holdout partition bounded between 1,000 and 10,000 rows. + + Args: + n_samples: Total number of raw rows in the input dataset. + + Returns: + Number of sample rows to isolate for global holdout validation. + """ + return min(10000, max(1000, int(0.05 * n_samples))) + def _get_n_features_to_add( self, n_features_requested: Union[int, str, None], n_cols: int @@ -1434,7 +1461,30 @@ def _generate_ensemble( if self.max_num_rows is not None: n_rows = min(n_rows, self.max_num_rows) - if n_rows < self.X_.shape[0]: + self.holdout_indices_ = None + if self.reserve_holdout_set: + n_val = self._compute_holdout_val_size(self.X_.shape[0]) + if n_val > 0.20 * self.X_.shape[0]: + raise ValueError( + f"Dataset size N ({self.X_.shape[0]}) is too small to reserve" + f" {n_val} holdout validation rows (would reduce training pool by >" + " 20%). To fix this, either increase max_num_rows to be >= N (or" + " set it to None) so all rows are used via cross-validation without" + " holdout deduction, or disable NNLS and calibration. Setting" + " max_num_rows >= 5000 will ensure there are always sufficient" + " rows to support a holdout set when needed." + ) + all_indices = list(range(self.X_.shape[0])) + val_indices_list = self.rng_.sample(all_indices, n_val) + val_set = set(val_indices_list) + train_pool = [i for i in all_indices if i not in val_set] + self.holdout_indices_ = np.array(val_indices_list) + actual_n_rows = min(len(train_pool), n_rows) + row_subsample_patterns = [ + np.array(self.rng_.sample(train_pool, actual_n_rows)) + for _ in range(self.n_estimators) + ] + elif n_rows < self.X_.shape[0]: row_subsample_patterns = [ np.array(self.rng_.sample(range(self.X_.shape[0]), n_rows)) for _ in range(self.n_estimators) @@ -1443,7 +1493,7 @@ def _generate_ensemble( row_subsample_patterns = [None] * self.n_estimators # 5. Combine into configurations - shuffle_shift_cat_configs = list( + config_tuples = list( zip( shuffle_patterns, shift_offsets, @@ -1451,14 +1501,14 @@ def _generate_ensemble( row_subsample_patterns, ) ) - self.rng_.shuffle(shuffle_shift_cat_configs) + self.rng_.shuffle(config_tuples) # 6. Assign Normalization Methods num_cycles = ( self.n_estimators + len(self.norm_methods_) - 1 ) // len(self.norm_methods_) norm_methods_for_estimators = (self.norm_methods_ * num_cycles)[: self.n_estimators] - full_configs = list(zip(norm_methods_for_estimators, shuffle_shift_cat_configs)) + full_configs = list(zip(norm_methods_for_estimators, config_tuples)) # Group by normalization method and separate components ensemble_configs: collections.OrderedDict = collections.OrderedDict() @@ -1535,33 +1585,76 @@ def transform_fold( X_test=None, train_fold=train_fold, val_fold=val_fold ) + def transform_holdout( + self, + ) -> Tuple[collections.OrderedDict, List[np.ndarray]]: + """Generate ensemble data views for global holdout validation evaluation. + + Relies on self.holdout_indices_ to define the validation set. + + Returns: + A tuple containing: + - data: Ordered dictionary mapping normalization string keys to + (X_ensemble, y_ensemble) arrays formatted identically to transform(). + - val_indices_list: List of absolute sample indices in the training set + corresponding to validation query rows for each ensemble member. + """ + assert self.holdout_indices_ is not None, "Holdout set not initialized." + return self._transform_features( + X_test=None, val_indices=self.holdout_indices_, + ) + def _transform_features( self, X_test: Optional[np.ndarray], train_fold: Optional[np.ndarray] = None, val_fold: Optional[np.ndarray] = None, + val_indices: Optional[np.ndarray] = None, ) -> Tuple[collections.OrderedDict, List[np.ndarray]]: """Shared helper to construct transformed feature and target dictionaries. Handles feature formatting, categorical value permutations, scaling, and - padding for both standard test inference (transform) and out-of-fold - cross-validation (transform_fold). + padding for standard test inference, out-of-fold cross-validation, and + global holdout validation evaluations. Either X_test alone, train_fold and + val_fold alone, or val_indices alone must be specified, corresponding to + each of these three cases. Args: X_test: 2D feature array of test queries of shape (n_samples, n_features). - If None, the evaluation test queries are taken from self.X_[val_fold] - during cross-validation. + If None, the evaluation test queries are taken from a validation or + hold-out portion of the training set self.X_. train_fold: Optional array of indices selecting in-context training rows during cross-validation. If None, all active training rows are used. val_fold: Optional array of indices selecting evaluation validation rows during cross-validation. + val_indices: Optional array of absolute row indices in self.X_ reserved + for global holdout validation. Returns: A tuple containing: - data: Ordered dictionary mapping normalization keys to transformed feature and target batches. - - val_indices_list: List of absolute validation row indices per config. + - val_indices_list: List of absolute validation row indices per ensemble + member, in the order of iterating over self.ensemble_configs_ in an + outer loop and self.ensemble_configs_[norm_method] in an inner loop. """ + assert (( + X_test is not None + and train_fold is None + and val_fold is None + and val_indices is None + ) or ( + X_test is None + and train_fold is not None + and val_fold is not None + and val_indices is None + ) or ( + X_test is None + and train_fold is None + and val_fold is None + and val_indices is not None + )) + y = self.y_ N = len(y) @@ -1570,9 +1663,9 @@ def _transform_features( max_features = 0 for ( norm_method, - shuffle_shift_cat_configs, + config_tuples, ) in self.ensemble_configs_.items(): - for shuffle_pattern, _, _, _ in shuffle_shift_cat_configs: + for shuffle_pattern, _, _, _ in config_tuples: max_features = max(max_features, len(shuffle_pattern)) data: collections.OrderedDict = collections.OrderedDict() @@ -1580,7 +1673,7 @@ def _transform_features( for ( norm_method, - shuffle_shift_cat_configs, + config_tuples, ) in self.ensemble_configs_.items(): preprocessor = self.preprocessors_[norm_method] X_ensemble = [] @@ -1591,16 +1684,26 @@ def _transform_features( shift_offset, cat_perm, row_sub_pattern, - ) in shuffle_shift_cat_configs: + ) in config_tuples: in_bag_idx = ( row_sub_pattern if row_sub_pattern is not None else np.arange(N) ) - if train_fold is not None and val_fold is not None: + if val_indices is not None: + # Global holdout validation: val_indices are direct absolute + # row indices in self.X_. + train_idx = in_bag_idx + val_idx = val_indices + val_indices_list.append(val_idx) + elif train_fold is not None and val_fold is not None: + # Standard CV fold evaluation: train_fold and val_fold are + # relative slices into in_bag_idx. train_idx = in_bag_idx[train_fold] val_idx = in_bag_idx[val_fold] val_indices_list.append(val_idx) else: + # Standard test inference: in_bag_idx contains all active + # training rows. train_idx = in_bag_idx val_idx = None @@ -1752,15 +1855,15 @@ def transform_context_only(self) -> collections.OrderedDict: max_features = 0 for ( norm_method, - shuffle_shift_cat_configs, + config_tuples, ) in self.ensemble_configs_.items(): - for shuffle_pattern, _, _, _ in shuffle_shift_cat_configs: + for shuffle_pattern, _, _, _ in config_tuples: max_features = max(max_features, len(shuffle_pattern)) data: collections.OrderedDict = collections.OrderedDict() for ( norm_method, - shuffle_shift_cat_configs, + config_tuples, ) in self.ensemble_configs_.items(): preprocessor = self.preprocessors_[norm_method] X_ensemble = [] @@ -1771,7 +1874,7 @@ def transform_context_only(self) -> collections.OrderedDict: shift_offset, cat_perm, row_sub_pattern, - ) in shuffle_shift_cat_configs: + ) in config_tuples: in_bag_idx = ( row_sub_pattern if row_sub_pattern is not None else np.arange(N) ) @@ -1830,15 +1933,15 @@ def transform_test_only(self, X: Any) -> collections.OrderedDict: max_features = 0 for ( norm_method, - shuffle_shift_cat_configs, + config_tuples, ) in self.ensemble_configs_.items(): - for shuffle_pattern, _, _, _ in shuffle_shift_cat_configs: + for shuffle_pattern, _, _, _ in config_tuples: max_features = max(max_features, len(shuffle_pattern)) data: collections.OrderedDict = collections.OrderedDict() for ( norm_method, - shuffle_shift_cat_configs, + config_tuples, ) in self.ensemble_configs_.items(): preprocessor = self.preprocessors_[norm_method] X_ensemble = [] @@ -1848,7 +1951,7 @@ def transform_test_only(self, X: Any) -> collections.OrderedDict: _, cat_perm, _, - ) in shuffle_shift_cat_configs: + ) in config_tuples: if cat_perm: X_test_to_use = X.copy() _apply_categorical_permutation(X_test_to_use, cat_perm) @@ -2375,10 +2478,6 @@ def __init__( self.keep_cache_on_device = keep_cache_on_device if self.average_logits and self.enable_nnls: raise ValueError("average_logits and enable_nnls cannot both be True.") - if self.max_num_rows is not None and self.enable_nnls: - raise ValueError( - "max_num_rows and enable_nnls cannot both be set at this time." - ) @classmethod def ensemble(cls, model: Any, **overrides: Any) -> "TabFMClassifier": @@ -2481,6 +2580,15 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": n_cat = 0 cat_features = list(range(n_cat)) + is_sampling = ( + self.max_num_rows is not None and self.max_num_rows < X.shape[0] + ) + needs_validation = self.enable_nnls or ( + self.active_calibration_method_ is not None + and self.active_calibration_method_ != "none" + ) + reserve_holdout = is_sampling and needs_validation + # Fit ensemble generator to create multiple dataset views self.ensemble_generator_ = EnsembleGenerator( n_estimators=self.n_estimators, @@ -2496,6 +2604,7 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": n_feature_crosses=self.n_feature_crosses, n_svd_features=self.n_svd_features, total_svd_pool=self.total_svd_pool, + reserve_holdout_set=reserve_holdout, ) self.ensemble_generator_.fit(X, y) @@ -2503,11 +2612,10 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": self._build_context_cache() oof_probs_fit = None - if self.enable_nnls or ( - self.active_calibration_method_ is not None - and self.active_calibration_method_ != "none" - ): - oof_probs = self.predict_oof_proba(cv=self.num_folds_for_cv) + y_orig_fit = None + y_fit = None + if needs_validation: + oof_probs, val_idx = self._predict_oof_proba(cv=self.num_folds_for_cv) val_idx = getattr(self, "oof_val_indices_", None) if val_idx is not None: oof_probs_fit = oof_probs[:, val_idx, :] @@ -2518,7 +2626,9 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": y_orig_fit = y_orig y_fit = y - if self.enable_nnls and oof_probs_fit is not None: + if self.enable_nnls: + assert oof_probs_fit is not None + assert y_orig_fit is not None n_classes = self.n_classes_ n_est, n_tr, _ = oof_probs_fit.shape @@ -2546,8 +2656,9 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": if ( self.active_calibration_method_ is not None and self.active_calibration_method_ != "none" - and oof_probs_fit is not None ): + assert oof_probs_fit is not None + assert y_fit is not None if self.enable_nnls: P = np.tensordot(self.ensemble_weights_, oof_probs_fit, axes=(0, 0)) else: @@ -2923,8 +3034,28 @@ def __getstate__(self): return state @jt.typed - def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K"]: - """Perform out-of-fold predictions on the training set for each ensemble member.""" + def _predict_oof_proba( + self, cv: int = 5 + ) -> tuple[jt.Float[jt.ArrayT, "E N K"], Optional[jt.Int[jt.ArrayT, "V"]]]: + """Perform out-of-fold predictions on some or all training samples. + + Predictions are either made using cross-validation, a single train/val + split, or a pre-computed holdout set. A holdout set is used when the + ensemble generator has a `holdout_indices` attribute. Otherwise, the + function decides whether to use cross-validation or a single train/val + split based on the number of rows and the `min_rows_for_single_val_split` + attribute. Regardless of the method, the output is structured such that (1) + there are predictions for every ensemble member on a shared set of rows and + (2) every prediction is generated without leakage. + + Args: + cv: Number of cross-validation folds to use for data splitting. + + Returns: + Tuple of (1) Out-of-fold probability predictions of shape (n_estimators, + n_samples, n_classes) and (2) the indices of the rows that have + predictions or None if all rows have predictions. + """ check_is_fitted(self) N = self.ensemble_generator_.X_.shape[0] @@ -2939,32 +3070,40 @@ def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K" all_configs.append((norm_method, config)) n_estimators = len(all_configs) - kf = KFold(n_splits=cv, shuffle=True, random_state=self.random_state) - # Take the number of rows from the first config if max_num_rows is given, - # or use N otherwise. - _, _, _, first_row_sub_pattern = all_configs[0][1] - n_rows = ( - len(first_row_sub_pattern) if first_row_sub_pattern is not None else N - ) - folds_base = list(kf.split(np.arange(n_rows))) + outputs_oof = np.zeros((n_estimators, N, n_classes)) + oof_val_indices = None - if ( - getattr(self, "min_rows_for_single_val_split", 0) > 0 - and len(folds_base[0][1]) >= self.min_rows_for_single_val_split - ): - folds_to_run = folds_base[:1] + if self.ensemble_generator_.holdout_indices is not None: + oof_val_indices = self.ensemble_generator_.holdout_indices + folds_data = [self.ensemble_generator_.transform_holdout()] else: - folds_to_run = folds_base + kf = KFold(n_splits=cv, shuffle=True, random_state=self.random_state) + # Take the number of rows from the first config if max_num_rows is given, + # or use N otherwise. + _, _, _, first_row_sub_pattern = all_configs[0][1] + n_rows = ( + len(first_row_sub_pattern) if first_row_sub_pattern is not None else N + ) + folds_base = list(kf.split(np.arange(n_rows))) - outputs_oof = np.zeros((n_estimators, N, n_classes)) + if ( + getattr(self, "min_rows_for_single_val_split", 0) > 0 + and len(folds_base[0][1]) >= self.min_rows_for_single_val_split + ): + folds_to_run = folds_base[:1] + else: + folds_to_run = folds_base - self.oof_val_indices_ = None - for fold_idx, (train_fold, val_fold) in enumerate(folds_to_run): - data_fold, val_indices_list = self.ensemble_generator_.transform_fold( - train_fold, val_fold - ) - if fold_idx == 0 and len(folds_to_run) == 1: - self.oof_val_indices_ = val_indices_list[0] + folds_data = [] + for fold_idx, (train_fold, val_fold) in enumerate(folds_to_run): + data_fold, val_indices_list = self.ensemble_generator_.transform_fold( + train_fold, val_fold + ) + if fold_idx == 0 and len(folds_to_run) == 1: + oof_val_indices = val_indices_list[0] + folds_data.append((data_fold, val_indices_list)) + + for data_fold, val_indices_list in folds_data: ( Xs_batch, ys_batch, @@ -2989,7 +3128,7 @@ def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K" ) outputs_oof[i, val_indices_list[i]] = out_i - return outputs_oof + return outputs_oof, oof_val_indices @jt.typed def _fit_calibration( @@ -3362,10 +3501,6 @@ def __init__( self.cache_context = cache_context self.maybe_quantize_kv_cache = maybe_quantize_kv_cache self.keep_cache_on_device = keep_cache_on_device - if self.max_num_rows is not None and self.enable_nnls: - raise ValueError( - "max_num_rows and enable_nnls cannot both be set at this time." - ) @classmethod def ensemble(cls, model: Any, **overrides: Any) -> "TabFMRegressor": @@ -3436,6 +3571,12 @@ def fit(self, X: Any, y: Any) -> "TabFMRegressor": self.y_scaler_.set_output(transform="default") y = self.y_scaler_.fit_transform(y.reshape(-1, 1)).flatten() + is_sampling = ( + self.max_num_rows is not None and self.max_num_rows < X.shape[0] + ) + needs_validation = self.enable_nnls + reserve_holdout = is_sampling and needs_validation + # Fit ensemble generator to create multiple dataset views self.ensemble_generator_ = EnsembleGenerator( n_estimators=self.n_estimators, norm_methods=self.norm_methods or ["none", "power"], @@ -3451,6 +3592,7 @@ def fit(self, X: Any, y: Any) -> "TabFMRegressor": n_feature_crosses=self.n_feature_crosses, n_svd_features=self.n_svd_features, total_svd_pool=self.total_svd_pool, + reserve_holdout_set=reserve_holdout, ) self.ensemble_generator_.fit(X, y) @@ -3458,10 +3600,9 @@ def fit(self, X: Any, y: Any) -> "TabFMRegressor": self._build_context_cache() if self.enable_nnls: - self.y_oof_scaled_ = self._compute_oof_preds_scaled( + self.y_oof_scaled_, val_idx = self._compute_oof_preds_scaled( cv=self.num_folds_for_cv ) - val_idx = getattr(self, "oof_val_indices_", None) y_oof_scaled_fit = ( self.y_oof_scaled_[:, val_idx] if val_idx is not None @@ -3822,11 +3963,30 @@ def predict_oof(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N"]: return outputs_oof + @jt.typed def _compute_oof_preds_scaled( self, cv: int = 5, - ) -> np.ndarray: - """Helper to compute OOF predictions in the scaled space.""" + ) -> tuple[jt.Float[jt.ArrayT, "E N"], Optional[jt.Int[jt.ArrayT, "V"]]]: + """Perform out-of-fold predictions on some or all training samples. + + Predictions are either made using cross-validation, a single train/val + split, or a pre-computed holdout set. A holdout set is used when the + ensemble generator has a `holdout_indices` attribute. Otherwise, the + function decides whether to use cross-validation or a single train/val + split based on the number of rows and the `min_rows_for_single_val_split` + attribute. Regardless of the method, the output is structured such that (1) + there are predictions for every ensemble member on a shared set of rows and + (2) every prediction is generated without leakage. + + Args: + cv: Number of cross-validation folds to use for data splitting. + + Returns: + Tuple of (1) Out-of-fold scaled predictions of shape (n_estimators, + n_samples, n_classes) and (2) the indices of the rows that have + predictions or None if all rows have predictions. + """ check_is_fitted(self) N = self.ensemble_generator_.X_.shape[0] @@ -3840,32 +4000,40 @@ def _compute_oof_preds_scaled( all_configs.append((norm_method, config)) n_estimators = len(all_configs) - kf = KFold(n_splits=cv, shuffle=True, random_state=self.random_state) - # Take the number of rows from the first config if max_num_rows is given, - # or use N otherwise. - _, _, _, first_row_sub_pattern = all_configs[0][1] - n_rows = ( - len(first_row_sub_pattern) if first_row_sub_pattern is not None else N - ) - folds_base = list(kf.split(np.arange(n_rows))) - outputs_oof = np.zeros((n_estimators, N)) + oof_val_indices = None - if ( - getattr(self, "min_rows_for_single_val_split", 0) > 0 - and len(folds_base[0][1]) >= self.min_rows_for_single_val_split - ): - folds_to_run = folds_base[:1] + if self.ensemble_generator_.holdout_indices is not None: + oof_val_indices = self.ensemble_generator_.holdout_indices + folds_data = [self.ensemble_generator_.transform_holdout()] else: - folds_to_run = folds_base - - self.oof_val_indices_ = None - for fold_idx, (train_fold, val_fold) in enumerate(folds_to_run): - data_fold, val_indices_list = self.ensemble_generator_.transform_fold( - train_fold, val_fold + kf = KFold(n_splits=cv, shuffle=True, random_state=self.random_state) + # Take the number of rows from the first config if max_num_rows is given, + # or use N otherwise. + _, _, _, first_row_sub_pattern = all_configs[0][1] + n_rows = ( + len(first_row_sub_pattern) if first_row_sub_pattern is not None else N ) - if fold_idx == 0 and len(folds_to_run) == 1: - self.oof_val_indices_ = val_indices_list[0] + folds_base = list(kf.split(np.arange(n_rows))) + + if ( + getattr(self, "min_rows_for_single_val_split", 0) > 0 + and len(folds_base[0][1]) >= self.min_rows_for_single_val_split + ): + folds_to_run = folds_base[:1] + else: + folds_to_run = folds_base + + folds_data = [] + for fold_idx, (train_fold, val_fold) in enumerate(folds_to_run): + data_fold, val_indices_list = self.ensemble_generator_.transform_fold( + train_fold, val_fold + ) + if fold_idx == 0 and len(folds_to_run) == 1: + oof_val_indices = val_indices_list[0] + folds_data.append((data_fold, val_indices_list)) + + for data_fold, val_indices_list in folds_data: ( Xs_batch, ys_batch, @@ -3884,7 +4052,7 @@ def _compute_oof_preds_scaled( for i in range(n_estimators): outputs_oof[i, val_indices_list[i]] = preds[i] - return outputs_oof + return outputs_oof, oof_val_indices @jt.typed def _predict_internal(self, X: Any) -> jt.Float[Array | np.ndarray, "E T"]: diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py index 2f3a299..e931d02 100644 --- a/tabfm/src/classifier_and_regressor_test.py +++ b/tabfm/src/classifier_and_regressor_test.py @@ -313,7 +313,7 @@ def mock_forward(Xs_batch, ys_batch, _cat_masks_batch=None, **_kwargs): with mock.patch.object( classifier, "_batch_forward", side_effect=mock_forward ) as mock_batch_forward: - oof_preds = classifier.predict_oof_proba(cv=2) + oof_preds, _ = classifier._predict_oof_proba(cv=2) self.assertEqual(oof_preds.shape, (2, 10, 2)) self.assertTrue(mock_batch_forward.called) @@ -468,7 +468,7 @@ def mock_forward(Xs_batch, ys_batch, _cat_masks_batch=None, **_kwargs): with mock.patch.object( regressor, "_batch_forward", side_effect=mock_forward ) as mock_batch_forward: - oof_preds = regressor.predict_oof(cv=2) + oof_preds, _ = regressor._compute_oof_preds_scaled(cv=2) self.assertEqual(oof_preds.shape, (2, 10)) self.assertTrue(mock_batch_forward.called) @@ -507,7 +507,7 @@ def mock_forward(Xs_batch, ys_batch, _cat_masks_batch=None, **_kwargs): with mock.patch.object( classifier, "_batch_forward", side_effect=mock_forward ) as mock_batch_forward: - oof_preds = classifier.predict_oof_proba(cv=2) + oof_preds, _ = classifier._predict_oof_proba(cv=2) self.assertEqual(oof_preds.shape, (2, 10, 2)) self.assertTrue(mock_batch_forward.called) @@ -701,7 +701,7 @@ def test_classifier_predict_oof_proba_raises_on_regression_model(self): classifier.fit(X, np.array([0, 1] * 5)) with self.assertRaisesRegex(ValueError, "fit on 2 classes"): - classifier.predict_oof_proba(cv=2) + classifier._predict_oof_proba(cv=2) @unittest.skipUnless(HAS_JAX, "JAX is required") @@ -738,7 +738,7 @@ def test_calibration_binary(self): ) with mock.patch.object( - classifier, "predict_oof_proba", return_value=mock_oof + classifier, _predict_oof_proba", return_value=(mock_oof, None) ): classifier.fit(X, y) @@ -783,7 +783,7 @@ def test_calibration_multiclass(self): ) with mock.patch.object( - classifier, "predict_oof_proba", return_value=mock_oof + classifier, "_predict_oof_proba", return_value=(mock_oof, None) ): classifier.fit(X, y) @@ -857,7 +857,7 @@ def test_classifier_nnls_blending(self): mock_oof = np.random.dirichlet([1, 1], size=(2, 10)) with mock.patch.object( - classifier, "predict_oof_proba", return_value=mock_oof + classifier, "_predict_oof_proba", return_value=(mock_oof, None) ): classifier.fit(X, y) @@ -890,7 +890,7 @@ def test_classifier_nnls_and_calibration(self): mock_oof = np.random.dirichlet([1, 1], size=(2, 50)) with mock.patch.object( - classifier, "predict_oof_proba", return_value=mock_oof + classifier, "_predict_oof_proba", return_value=(mock_oof, None) ): with mock.patch.object(classifier, "_fit_calibration") as mock_fit_cal: classifier.fit(X, y) @@ -918,7 +918,7 @@ def test_classifier_nnls_and_calibration_small_dataset(self): mock_oof = np.random.dirichlet([1, 1], size=(2, 10)) with mock.patch.object( - classifier, "predict_oof_proba", return_value=mock_oof + classifier, "_predict_oof_proba", return_value=(mock_oof, None) ): classifier.fit(X, y) @@ -945,10 +945,11 @@ def test_max_num_rows(self): mock_oof = np.random.dirichlet([1, 1], size=(2, 10)) with mock.patch.object( - classifier, "predict_oof_proba", return_value=mock_oof + classifier, "_predict_oof_proba", return_value=(mock_oof, None) ): classifier.fit(X, y) + self.assertIsNone(classifier.ensemble_generator_.holdout_indices_) self.assertEqual(classifier.ensemble_generator_.X_.shape[0], 10) patterns_none = classifier.ensemble_generator_.row_subsample_patterns_[ "none" @@ -998,6 +999,89 @@ def test_min_rows_for_single_val_split_regressor(self): self.assertEqual(mock_forward.call_count, 1) + def test_max_num_rows_with_enable_nnls_too_small_raises(self): + classifier = TabFMClassifier( + model=self.model, + config=self.config, + n_estimators=2, + max_num_rows=2000, + average_logits=False, + enable_nnls=True, + ) + X = np.random.rand(4000, 3) + y = np.array([0, 1] * 2000) + with self.assertRaisesRegex(ValueError, "reduce training pool by > 20%"): + classifier.fit(X, y) + + def test_max_num_rows_with_enable_nnls_classifier(self): + classifier = TabFMClassifier( + model=self.model, + config=self.config, + n_estimators=2, + max_num_rows=5000, + enable_nnls=True, + average_logits=False, + num_folds_for_cv=5, + ) + X = np.random.rand(10000, 3) + y = np.array([0, 1] * 5000) + + # 5% of 10000 is 500, max(1000, 500) = 1000 validation holdout rows. + # Left for training pool: 10000 - 1000 = 9000. + with mock.patch.object( + classifier, "_batch_forward", return_value=np.zeros((2, 1000, 2)) + ) as mock_forward: + classifier.fit(X, y) + + self.assertIsNotNone(classifier.ensemble_generator_.holdout_indices_) + self.assertLen(classifier.ensemble_generator_.holdout_indices_, 1000) + self.assertEqual(mock_forward.call_count, 1) + + def test_max_num_rows_with_enable_nnls_regressor(self): + regressor = TabFMRegressor( + model=self.model, + config=argparse.Namespace(loss="rmse"), + n_estimators=2, + max_num_rows=5000, + enable_nnls=True, + num_folds_for_cv=5, + ) + X = np.random.rand(40000, 3) + y = np.random.rand(40000) + + # 5% of 40000 is 2000 (true 5% regime between 1k floor and 10k cap). + with mock.patch.object( + regressor, "_batch_forward", return_value=np.zeros((2, 2000, 1)) + ) as mock_forward: + regressor.fit(X, y) + + self.assertIsNotNone(regressor.ensemble_generator_.holdout_indices_) + self.assertLen(regressor.ensemble_generator_.holdout_indices_, 2000) + self.assertEqual(mock_forward.call_count, 1) + + def test_max_num_rows_greater_than_N_with_enable_nnls_uses_cv(self): + classifier = TabFMClassifier( + model=self.model, + config=self.config, + n_estimators=2, + max_num_rows=6000, + enable_nnls=True, + average_logits=False, + num_folds_for_cv=5, + ) + X = np.random.rand(5500, 3) + y = np.array([0, 1] * 2750) + + # With max_num_rows=6000 >= N=5500, no holdout set is partitioned and + # evaluation runs standard 5-fold CV (1100 validation rows per fold). + with mock.patch.object( + classifier, "_batch_forward", return_value=np.zeros((2, 1100, 2)) + ) as mock_forward: + classifier.fit(X, y) + + self.assertIsNone(classifier.ensemble_generator_.holdout_indices_) + self.assertEqual(mock_forward.call_count, 5) + class EnsemblePresetTest(absltest.TestCase): From f9015af5a93bf93e95fe9afa5e80c60dedaded56 Mon Sep 17 00:00:00 2001 From: tamannarayan Date: Tue, 28 Jul 2026 12:18:45 -0400 Subject: [PATCH 2/4] Fix typing and remove dead code. --- tabfm/src/classifier_and_regressor.py | 21 ++++++--------------- tabfm/src/classifier_and_regressor_test.py | 2 +- 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index e263d84..77a88c4 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/tabfm/src/classifier_and_regressor.py @@ -3036,7 +3036,9 @@ def __getstate__(self): @jt.typed def _predict_oof_proba( self, cv: int = 5 - ) -> tuple[jt.Float[jt.ArrayT, "E N K"], Optional[jt.Int[jt.ArrayT, "V"]]]: + ) -> tuple[ + jt.Float[jt.Array | np.ndarray, "E N K"], + Optional[jt.Int[jt.Array | np.ndarray, "V"]]]: """Perform out-of-fold predictions on some or all training samples. Predictions are either made using cross-validation, a single train/val @@ -3949,25 +3951,14 @@ def __getstate__(self): state.pop(attr, None) return state - @jt.typed - def predict_oof(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N"]: - """Perform out-of-fold predictions on the training set for each ensemble member.""" - check_is_fitted(self) - - outputs_oof_scaled = self._compute_oof_preds_scaled(cv=cv) - n_estimators, N = outputs_oof_scaled.shape - - outputs_oof = np.zeros((n_estimators, N)) - for i in range(n_estimators): - outputs_oof[i, :] = self._inverse_transform_y(outputs_oof_scaled[i]) - - return outputs_oof @jt.typed def _compute_oof_preds_scaled( self, cv: int = 5, - ) -> tuple[jt.Float[jt.ArrayT, "E N"], Optional[jt.Int[jt.ArrayT, "V"]]]: + ) -> tuple[ + jt.Float[jt.Array | np.ndarray, "E N"], + Optional[jt.Int[jt.Array | np.ndarray, "V"]]]: """Perform out-of-fold predictions on some or all training samples. Predictions are either made using cross-validation, a single train/val diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py index e931d02..4de02eb 100644 --- a/tabfm/src/classifier_and_regressor_test.py +++ b/tabfm/src/classifier_and_regressor_test.py @@ -685,7 +685,7 @@ def test_regressor_predict_oof_raises_on_classification_model(self): regressor.fit(X, np.random.rand(10)) with self.assertRaisesRegex(ValueError, "model_type='regression'"): - regressor.predict_oof(cv=2) + regressor._compute_oof_preds_scaled(cv=2) def test_classifier_predict_proba_raises_on_regression_model(self): classifier = TabFMClassifier(model=self._tiny_model("rmse"), n_estimators=2) From ec931b8ef85972203a7de31a9a295987c69595e0 Mon Sep 17 00:00:00 2001 From: tamannarayan Date: Tue, 28 Jul 2026 12:32:15 -0400 Subject: [PATCH 3/4] Fix typo --- tabfm/src/classifier_and_regressor_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py index 4de02eb..eb221f7 100644 --- a/tabfm/src/classifier_and_regressor_test.py +++ b/tabfm/src/classifier_and_regressor_test.py @@ -738,7 +738,7 @@ def test_calibration_binary(self): ) with mock.patch.object( - classifier, _predict_oof_proba", return_value=(mock_oof, None) + classifier, "_predict_oof_proba", return_value=(mock_oof, None) ): classifier.fit(X, y) From 01004165b72a86305886ba06a67e7aae7c733328 Mon Sep 17 00:00:00 2001 From: tamannarayan Date: Tue, 28 Jul 2026 13:05:43 -0400 Subject: [PATCH 4/4] Test fixes. --- tabfm/src/classifier_and_regressor_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py index eb221f7..419c983 100644 --- a/tabfm/src/classifier_and_regressor_test.py +++ b/tabfm/src/classifier_and_regressor_test.py @@ -1002,7 +1002,6 @@ def test_min_rows_for_single_val_split_regressor(self): def test_max_num_rows_with_enable_nnls_too_small_raises(self): classifier = TabFMClassifier( model=self.model, - config=self.config, n_estimators=2, max_num_rows=2000, average_logits=False, @@ -1016,7 +1015,6 @@ def test_max_num_rows_with_enable_nnls_too_small_raises(self): def test_max_num_rows_with_enable_nnls_classifier(self): classifier = TabFMClassifier( model=self.model, - config=self.config, n_estimators=2, max_num_rows=5000, enable_nnls=True, @@ -1040,7 +1038,6 @@ def test_max_num_rows_with_enable_nnls_classifier(self): def test_max_num_rows_with_enable_nnls_regressor(self): regressor = TabFMRegressor( model=self.model, - config=argparse.Namespace(loss="rmse"), n_estimators=2, max_num_rows=5000, enable_nnls=True, @@ -1062,7 +1059,6 @@ def test_max_num_rows_with_enable_nnls_regressor(self): def test_max_num_rows_greater_than_N_with_enable_nnls_uses_cv(self): classifier = TabFMClassifier( model=self.model, - config=self.config, n_estimators=2, max_num_rows=6000, enable_nnls=True,