From f649eae5073f1e786c17c0d5106042f8976cfef1 Mon Sep 17 00:00:00 2001 From: Pavlos Nicolaou Date: Tue, 21 Jul 2026 22:13:02 +0300 Subject: [PATCH] Fix four open bugs: MPS float64 crash, OOF calibration with max_num_rows, sklearn transform_output leak, missing safetensors dep - Cast float64 targets to float32 before the device move in _predict_step_pytorch and the context-cache prefill: MPS rejects float64 tensors at transfer time, so the existing guard was unreachable on Apple Silicon (#68). - Classifier calibration with max_num_rows: average out-of-fold probabilities per row over the members that actually predicted it, instead of slicing all members by member 0's validation indices, which mixed all-zero rows into the calibration fit (#55). - Run public fit/predict entry points under sklearn.config_context(transform_output="default") so a global sklearn.set_config(transform_output="pandas") cannot break the numpy-based internal pipeline (#58). - Declare safetensors in the pytorch extra: the published HF checkpoint is safetensors-only, and loading it without the package fails with NameError inside huggingface_hub (#56). --- CHANGELOG.md | 18 +++ pyproject.toml | 1 + requirements.txt | 2 + tabfm/src/classifier_and_regressor.py | 62 +++++++++-- .../classifier_and_regressor_pytorch_test.py | 105 ++++++++++++++++++ 5 files changed, 176 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ca8ca44..7227eb4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,23 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`): --> +## [Unreleased] + +### Fixed + +* PyTorch backend: float64 targets are now cast to float32 before the device + move (in both the predict step and context prefill), fixing a crash on Apple + MPS, which rejects float64 tensors at transfer time. (#68) +* Classifier calibration with `max_num_rows`: out-of-fold probabilities are now + averaged per row over the ensemble members that actually predicted that row, + instead of mixing in all-zero rows from members whose row subsample did not + cover it. (#55) +* `TabFMClassifier`/`TabFMRegressor` are no longer affected by a global + `sklearn.set_config(transform_output="pandas")`; internal transformers always + produce numpy arrays. (#58) +* The `pytorch` extra now declares the `safetensors` dependency required to + load the Hugging Face checkpoint. (#56) + ## [1.0.1] - 2026-07-09 ### Fixed @@ -59,5 +76,6 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`): * Initial release +[Unreleased]: https://github.com/google-research/tabfm/compare/v1.0.1...HEAD [1.0.1]: https://github.com/google-research/tabfm/compare/v1.0.0...v1.0.1 [1.0.0]: https://github.com/google-research/tabfm/releases/tag/v1.0.0 diff --git a/pyproject.toml b/pyproject.toml index 784274d..60532c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ jax = [ "orbax-checkpoint", ] pytorch = [ + "safetensors", "torch", ] diff --git a/requirements.txt b/requirements.txt index 9a5e4e4..90e3978 100644 --- a/requirements.txt +++ b/requirements.txt @@ -151,6 +151,8 @@ rich==13.7.1 # via # flax # typer +safetensors==0.7.0 + # via tabfm (pyproject.toml) scikit-learn==1.6.0 # via tabfm (pyproject.toml) scipy==1.17.1 diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index 4a0ebee..a024d43 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/tabfm/src/classifier_and_regressor.py @@ -27,6 +27,7 @@ """ import collections +import functools import itertools import math import random @@ -56,6 +57,7 @@ import pandas as pd import scipy.optimize as opt import scipy.special +import sklearn from sklearn.base import BaseEstimator from sklearn.base import ClassifierMixin from sklearn.base import RegressorMixin @@ -366,6 +368,22 @@ def transform(self, X: Any) -> np.ndarray: return X_datetime.values +def _with_default_transform_output(method): + """Runs ``method`` under sklearn's default transform-output config. + + User code may globally request pandas transform output via + ``sklearn.set_config(transform_output="pandas")``; the numpy-based internal + pipeline must not be affected by that global setting. + """ + + @functools.wraps(method) + def wrapper(*args, **kwargs): + with sklearn.config_context(transform_output="default"): + return method(*args, **kwargs) + + return wrapper + + class TransformToNumerical(TransformerMixin, BaseEstimator): """Transforms non-numerical data in a DataFrame to numerical representations. @@ -1977,9 +1995,12 @@ def _predict_step_pytorch( device = next(model.parameters()).device X_t = torch.from_numpy(X_batch).to(device, dtype=torch.float32) - y_t = torch.from_numpy(y_batch).to(device) + # Downcast float64 targets host-side: MPS rejects float64 tensors at + # transfer time, so the cast must happen before the device move. + y_t = torch.from_numpy(y_batch) if y_t.dtype == torch.float64: y_t = y_t.to(torch.float32) + y_t = y_t.to(device) batch_size = X_batch.shape[0] train_size_t = torch.full( @@ -2064,9 +2085,12 @@ def _build_context_cache_pytorch( Xs_split, ys_split, cat_masks_split, ds_split ): X_t = torch.from_numpy(X_batch).to(device, dtype=torch.float32) - y_t = torch.from_numpy(y_batch).to(device) + # Downcast float64 targets host-side: MPS rejects float64 tensors at + # transfer time, so the cast must happen before the device move. + y_t = torch.from_numpy(y_batch) if y_t.dtype == torch.float64: y_t = y_t.to(torch.float32) + y_t = y_t.to(device) cat_mask_t = torch.from_numpy(cat_mask_batch).to(device) d_t = torch.from_numpy(ds_batch).to(device) _, cache = model.prefill(X_t, y_t, cat_mask=cat_mask_t, d=d_t) @@ -2417,6 +2441,7 @@ def __sklearn_tags__(self): tags.non_deterministic = True return tags + @_with_default_transform_output def fit(self, X: Any, y: Any) -> "TabFMClassifier": """Fit the classifier to training data. @@ -2508,15 +2533,17 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": and self.active_calibration_method_ != "none" ): oof_probs = 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, :] - y_orig_fit = y_orig[val_idx] - y_fit = y[val_idx] - else: - oof_probs_fit = oof_probs - y_orig_fit = y_orig - y_fit = y + # With max_num_rows each ensemble member draws its own row subsample, so + # a row may carry OOF predictions from only a subset of members. Keep the + # rows at least one member predicted (with per-member coverage recorded + # in oof_counts_fit) instead of slicing all members by member 0's + # validation indices, which would mix in all-zero rows for the others. + oof_counts = self.oof_pred_mask_.sum(axis=0) + oof_rows = np.flatnonzero(oof_counts > 0) + oof_probs_fit = oof_probs[:, oof_rows, :] + oof_counts_fit = oof_counts[oof_rows] + y_orig_fit = y_orig[oof_rows] + y_fit = y[oof_rows] if self.enable_nnls and oof_probs_fit is not None: n_classes = self.n_classes_ @@ -2551,7 +2578,9 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": if self.enable_nnls: P = np.tensordot(self.ensemble_weights_, oof_probs_fit, axes=(0, 0)) else: - P = np.mean(oof_probs_fit, axis=0) + # Average each row over the members that predicted it; unpredicted + # member entries are zero and must not drag the mean toward zero. + P = oof_probs_fit.sum(axis=0) / oof_counts_fit[:, None] assert P.shape == (len(y_fit), self.n_classes_), ( f"Expected calibration input shape {(len(y_fit), self.n_classes_)}," f" got {P.shape}" @@ -2922,6 +2951,7 @@ def __getstate__(self): state.pop(attr, None) return state + @_with_default_transform_output @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.""" @@ -2957,6 +2987,7 @@ def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K" folds_to_run = folds_base outputs_oof = np.zeros((n_estimators, N, n_classes)) + oof_mask = np.zeros((n_estimators, N), dtype=bool) self.oof_val_indices_ = None for fold_idx, (train_fold, val_fold) in enumerate(folds_to_run): @@ -2988,7 +3019,9 @@ def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K" out_i, axis=-1, temperature=self.softmax_temperature ) outputs_oof[i, val_indices_list[i]] = out_i + oof_mask[i, val_indices_list[i]] = True + self.oof_pred_mask_ = oof_mask return outputs_oof @jt.typed @@ -3165,6 +3198,7 @@ def _process_logits(self, logits_all: np.ndarray): return probs + @_with_default_transform_output @jt.typed def predict_proba(self, X: Any) -> jt.Float[Array | np.ndarray, "T K"]: """Predict class probabilities for test samples. @@ -3191,6 +3225,7 @@ def predict_proba(self, X: Any) -> jt.Float[Array | np.ndarray, "T K"]: logits = self._predict_proba_internal(X) return self._process_logits(logits) + @_with_default_transform_output @jt.typed def predict(self, X: Any) -> np.ndarray: """Predict class labels for test samples. @@ -3395,6 +3430,7 @@ def _more_tags(self): """Mark regressor as non-deterministic to bypass certain sklearn tests.""" return dict(non_deterministic=True) + @_with_default_transform_output def fit(self, X: Any, y: Any) -> "TabFMRegressor": """Fit the regressor to training data. @@ -3804,6 +3840,7 @@ def __getstate__(self): state.pop(attr, None) return state + @_with_default_transform_output @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.""" @@ -3923,6 +3960,7 @@ def _combine_predictions( avg_predictions = np.mean(predictions_scaled, axis=0) return self._inverse_transform_y(avg_predictions) + @_with_default_transform_output @jt.typed def predict(self, X: Any) -> jt.Float[Array | np.ndarray, "T"]: """Predict regression target for test samples. diff --git a/tabfm/src/classifier_and_regressor_pytorch_test.py b/tabfm/src/classifier_and_regressor_pytorch_test.py index 4ea1df0..246c592 100644 --- a/tabfm/src/classifier_and_regressor_pytorch_test.py +++ b/tabfm/src/classifier_and_regressor_pytorch_test.py @@ -17,12 +17,45 @@ from unittest import mock import numpy as np import pandas as pd +import sklearn import torch from tabfm.src.pytorch import model as pytorch_model +from tabfm.src import classifier_and_regressor from tabfm.src.classifier_and_regressor import TabFMClassifier, TabFMRegressor +def _small_model(is_classifier: bool, max_classes: int) -> pytorch_model.TabFM: + return pytorch_model.TabFM( + embed_dim=8, + max_classes=max_classes, + col_num_blocks=1, + col_nhead=2, + col_num_inds=8, + row_num_blocks=1, + row_nhead=2, + row_num_cls=2, + icl_num_blocks=1, + icl_nhead=2, + ff_factor=2, + feature_group_size=2, + is_classifier=is_classifier, + ) + + +class _DTypeCaptureModel(torch.nn.Module): + """Records the dtype of the target tensor handed to the model.""" + + def __init__(self): + super().__init__() + self.param = torch.nn.Parameter(torch.zeros(1)) + self.seen_y_dtype = None + + def forward(self, X, y, train_size, cat_mask=None, d=None): + self.seen_y_dtype = y.dtype + return torch.zeros(X.shape[0], X.shape[1], 1) + + class PyTorchClassifierRegressorTest(unittest.TestCase): def test_classifier_fit_predict(self): @@ -136,6 +169,78 @@ def test_regressor_fit_predict(self): self.assertEqual(preds_cached.shape, (10,)) np.testing.assert_allclose(preds_cached, preds, rtol=1e-5, atol=1e-6) + def test_predict_step_casts_float64_targets_before_device_move(self): + model = _DTypeCaptureModel() + classifier_and_regressor._predict_step_pytorch( + model, + np.random.rand(2, 6, 3).astype(np.float32), + np.random.rand(2, 4), # float64, numpy's default float dtype + 4, + None, + None, + ) + self.assertEqual(model.seen_y_dtype, torch.float32) + + def test_calibration_with_max_num_rows_uses_valid_probabilities(self): + np.random.seed(0) + model = _small_model(is_classifier=True, max_classes=3) + + orig_fit_calibration = TabFMClassifier._fit_calibration + captured = {} + + def spy(clf_self, P, y_fit): + captured["P"] = np.asarray(P) + return orig_fit_calibration(clf_self, P, y_fit) + + with mock.patch.object(TabFMClassifier, "_fit_calibration", spy): + clf = TabFMClassifier( + model=model, + n_estimators=3, + batch_size=3, + random_state=0, + max_num_rows=30, + min_rows_for_single_val_split=1, + binary_calibration_method="platt", + ) + X = np.random.rand(60, 3) + y = np.random.randint(0, 2, size=60) + clf.fit(X, y) + + # Each member subsamples rows independently, so the calibration inputs + # must be averaged per row over the members that predicted it — every + # row must still be a probability vector. + P = captured["P"] + self.assertGreater(P.shape[0], 0) + np.testing.assert_allclose(P.sum(axis=1), 1.0, rtol=1e-5) + + def test_fit_predict_with_sklearn_pandas_output(self): + np.random.seed(42) + sklearn.set_config(transform_output="pandas") + try: + reg = TabFMRegressor( + model=_small_model(is_classifier=False, max_classes=1), + n_estimators=2, + batch_size=2, + random_state=42, + ) + X = np.random.rand(10, 3) + y = np.random.rand(10) + reg.fit(X, y) + self.assertEqual(reg.predict(X).shape, (10,)) + + clf = TabFMClassifier( + model=_small_model(is_classifier=True, max_classes=3), + n_estimators=2, + batch_size=2, + random_state=42, + ) + X_df = pd.DataFrame(np.random.rand(10, 3), columns=["a", "b", "c"]) + y_cls = np.array([0, 1, 2, 0, 1, 2, 0, 1, 2, 0]) + clf.fit(X_df, y_cls) + self.assertEqual(clf.predict_proba(X_df).shape, (10, 3)) + finally: + sklearn.set_config(transform_output="default") + class PyTorchModelPickleTest(unittest.TestCase): """The PyTorch model must be picklable.