Fix open bugs: MPS float64 crash (#68), OOF calibration with max_num_rows (#55), sklearn transform_output leak (#58), safetensors dep (#56)

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

pyproject.toml

@@ -54,6 +54,7 @@ jax = [
     "orbax-checkpoint",
 ]
 pytorch = [
+    "safetensors",
     "torch",
 ]
 

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

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.

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.