From 79b0b6ba899b4c853417b8140753887133f639e3 Mon Sep 17 00:00:00 2001 From: Benjamin Jaeger Date: Wed, 23 Sep 2026 14:37:02 +0200 Subject: [PATCH 1/5] Give each regression member its own target pipeline [RES-2639] --- changelog/1198.changed.md | 1 + src/tabpfn/finetuning/data_util.py | 15 +- src/tabpfn/finetuning/finetuned_regressor.py | 16 +- src/tabpfn/preprocessing/target_transform.py | 72 +++++++++ src/tabpfn/regressor.py | 47 ++++-- src/tabpfn/utils.py | 15 +- tests/test_batched_prediction_shape_safety.py | 9 +- tests/test_downsample_correction.py | 9 +- tests/test_finetuning_regressor.py | 47 ++++++ .../test_target_transform.py | 151 ++++++++++++++++++ tests/test_regressor_interface.py | 63 ++++++++ 11 files changed, 416 insertions(+), 29 deletions(-) create mode 100644 changelog/1198.changed.md create mode 100644 src/tabpfn/preprocessing/target_transform.py create mode 100644 tests/test_preprocessing/test_target_transform.py diff --git a/changelog/1198.changed.md b/changelog/1198.changed.md new file mode 100644 index 000000000..5cb4273f1 --- /dev/null +++ b/changelog/1198.changed.md @@ -0,0 +1 @@ +Regression ensembles now normalize each member’s target using its own training rows, correcting predictions when row subsampling is enabled. Without row subsampling, predictions may differ slightly due to floating-point rounding. diff --git a/src/tabpfn/finetuning/data_util.py b/src/tabpfn/finetuning/data_util.py index d461b99da..944a8e7bc 100644 --- a/src/tabpfn/finetuning/data_util.py +++ b/src/tabpfn/finetuning/data_util.py @@ -66,6 +66,8 @@ class RegressorBatch: znorm_space_bardist: Bar distribution in z-normalized target space. X_query_raw: Original unprocessed test features. y_query_raw: Original unprocessed test targets. + y_train_mean: Training-target mean defining the shared aggregation frame. + y_train_std: Standard deviation of this split's training target. """ X_context: list[torch.Tensor] @@ -79,6 +81,8 @@ class RegressorBatch: znorm_space_bardist: FullSupportBarDistribution X_query_raw: torch.Tensor y_query_raw: torch.Tensor + y_train_mean: float = 0.0 + y_train_std: float = 1.0 Batch = TypeVar("Batch", ClassifierBatch, RegressorBatch) @@ -453,7 +457,6 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: train_std = eps y_test_standardized = (y_test_raw - train_mean) / train_std - y_train_standardized = (y_train_raw - train_mean) / train_std znorm_borders = znorm_space_bardist_.borders raw_space_bardist_ = ( FullSupportBarDistribution( @@ -463,9 +466,6 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: .float() .to(znorm_borders.device) ) - y_train = y_train_standardized - else: - y_train = y_train_raw num_columns = x_train_raw.shape[1] feature_schema = FeatureSchema.from_only_categorical_indices( @@ -480,7 +480,7 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: ) ensemble_members = ensemble_preprocessor.fit_transform_ensemble_members( X_train=x_train_raw, - y_train=y_train, + y_train=y_train_raw, ) X_trains_preprocessed = [m.X_train for m in ensemble_members] y_trains_preprocessed = [m.y_train for m in ensemble_members] @@ -535,6 +535,8 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: znorm_space_bardist=znorm_space_bardist_, X_query_raw=x_test_raw, y_query_raw=y_test_raw, + y_train_mean=float(train_mean), + y_train_std=float(train_std), ) return ClassifierBatch( @@ -700,6 +702,9 @@ def meta_dataset_collator( znorm_space_bardist=first_item.znorm_space_bardist, X_query_raw=_collate_tensor_field(batch, "X_query_raw", padding_val), y_query_raw=_collate_tensor_field(batch, "y_query_raw", padding_val), + # The target frame belongs to the bar distributions above. + y_train_mean=first_item.y_train_mean, + y_train_std=first_item.y_train_std, ) diff --git a/src/tabpfn/finetuning/finetuned_regressor.py b/src/tabpfn/finetuning/finetuned_regressor.py index b2d3449e5..f4f539fda 100644 --- a/src/tabpfn/finetuning/finetuned_regressor.py +++ b/src/tabpfn/finetuning/finetuned_regressor.py @@ -40,7 +40,7 @@ def _targets_in_estimator_space( y_Q: torch.Tensor, config: RegressorEnsembleConfig ) -> torch.Tensor: - """Map z-scored targets into the space the estimator predicts in. + """Map raw targets into the space the estimator predicts in. Members with a target transform (the default ensemble's "safepower") saw transformed context targets, so their loss must use transformed targets too. @@ -48,7 +48,7 @@ def _targets_in_estimator_space( if config.target_transform is None: return y_Q y_np = config.target_transform.transform(y_Q.cpu().numpy().reshape(-1, 1)) - return torch.as_tensor(y_np.ravel(), dtype=y_Q.dtype, device=y_Q.device) + return torch.as_tensor(y_np.ravel(), device=y_Q.device) def _compute_regression_loss( # noqa: C901 @@ -368,10 +368,13 @@ def _should_skip_batch(self, batch: RegressorBatch) -> bool: # type: ignore[ove @override def _setup_batch(self, batch: RegressorBatch) -> None: # type: ignore[override] - """Set up bar distribution for this batch.""" + """Set up bar distribution and target frame for this batch.""" self.finetuned_estimator_.raw_space_bardist_ = batch.raw_space_bardist self.finetuned_estimator_.bardist_ = batch.znorm_space_bardist self._bardist_loss = batch.znorm_space_bardist + # Preprocessed fitting needs the batch's target frame for border mapping. + self.finetuned_estimator_.y_train_mean_ = batch.y_train_mean + self.finetuned_estimator_.y_train_std_ = batch.y_train_std @override def _forward_with_loss(self, batch: RegressorBatch) -> torch.Tensor: # type: ignore[override] @@ -406,8 +409,11 @@ def _forward_with_loss(self, batch: RegressorBatch) -> torch.Tensor: # type: ig logits_BQL = logits_QBEL.permute(1, 2, 0, 3).reshape(B * E, Q, L) targets_BQ = torch.stack( - [_targets_in_estimator_space(y_query_batch[0], c[0]) for c in batch.configs] - ).to(self.device) + [ + _targets_in_estimator_space(batch.y_query_raw[0], c[0]) + for c in batch.configs + ] + ).to(device=self.device, dtype=y_query_batch.dtype) return _compute_regression_loss( logits_BQL=logits_BQL, diff --git a/src/tabpfn/preprocessing/target_transform.py b/src/tabpfn/preprocessing/target_transform.py new file mode 100644 index 000000000..53fc2d669 --- /dev/null +++ b/src/tabpfn/preprocessing/target_transform.py @@ -0,0 +1,72 @@ +# Copyright (c) Prior Labs GmbH 2026. + +"""Invertible regression target pipelines. + +Each member maps raw targets to model inputs and inversely maps borders to raw +units. The estimator maps those borders into a shared standardized space for +aggregation, avoiding float32 precision loss for targets with large offsets. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from sklearn.base import BaseEstimator, TransformerMixin +from sklearn.pipeline import Pipeline + +if TYPE_CHECKING: + from sklearn.base import TransformerMixin as Transformer + +STANDARDIZE_STEP = "standardize_target" +TARGET_TRANSFORM_STEP = "target_transform" + + +class StandardizeTarget(TransformerMixin, BaseEstimator): + """Standardize using fitted ``mean_`` and ``std_`` (population std + epsilon). + + Keep NumPy arithmetic rather than StandardScaler to preserve exact model + inputs when the member sees the full training target. + """ + + EPSILON = 1e-20 + """Prevent division by zero for constant targets.""" + + def fit(self, X: np.ndarray, y: np.ndarray | None = None) -> StandardizeTarget: + """Learn the mean and standard deviation of the target ``X``.""" + del y + self.mean_ = float(np.mean(X)) + self.std_ = float(np.std(X)) + self.EPSILON + return self + + def transform(self, X: np.ndarray) -> np.ndarray: + """Return the standardized target.""" + return (np.asarray(X) - self.mean_) / self.std_ + + def inverse_transform(self, X: np.ndarray) -> np.ndarray: + """Return ``X`` in the original units of the target.""" + return np.asarray(X) * self.std_ + self.mean_ + + +def make_target_transform(transform: Transformer | Pipeline | None) -> Pipeline: + """Standardize raw targets, then apply the optional preset. + + The inverse returns values in the target's original units. + """ + if transform is None: + return Pipeline(steps=[(STANDARDIZE_STEP, StandardizeTarget())]) + return Pipeline( + steps=[ + # The preset reshapes the standardized target. + (STANDARDIZE_STEP, StandardizeTarget()), + (TARGET_TRANSFORM_STEP, transform), + ], + ) + + +__all__ = [ + "STANDARDIZE_STEP", + "TARGET_TRANSFORM_STEP", + "StandardizeTarget", + "make_target_transform", +] diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index b47747137..a02c7638f 100644 --- a/src/tabpfn/regressor.py +++ b/src/tabpfn/regressor.py @@ -101,6 +101,10 @@ from tabpfn.preprocessing.steps import ( get_all_reshape_feature_distribution_preprocessors, ) +from tabpfn.preprocessing.target_transform import ( + StandardizeTarget, + make_target_transform, +) from tabpfn.preprocessing.text import TextTransformer from tabpfn.utils import ( DevicesSpecification, @@ -1017,15 +1021,17 @@ def _initialize_dataset_preprocessing( num_examples=y.shape[0], # Use length of validated y random_state=random_state, # Use the provided rng ) - target_preprocessors: list[TransformerMixin | Pipeline | None] = [] - for ( - y_target_preprocessor - ) in self.inference_config_.REGRESSION_Y_PREPROCESS_TRANSFORMS: - if y_target_preprocessor is not None: - preprocessor = possible_target_transforms[y_target_preprocessor] - else: - preprocessor = None - target_preprocessors.append(preprocessor) + # Each member fits its target pipeline on raw targets. + target_preprocessors: list[TransformerMixin | Pipeline | None] = [ + make_target_transform( + None + if y_target_preprocessor is None + else possible_target_transforms[y_target_preprocessor] + ) + for y_target_preprocessor in ( + self.inference_config_.REGRESSION_Y_PREPROCESS_TRANSFORMS + ) + ] preprocessor_configs = self.inference_config_.PREPROCESS_TRANSFORMS self.n_estimators_ = scale_n_estimators_for_feature_coverage( @@ -1365,11 +1371,9 @@ def fit(self, X: XType, y: YType) -> Self: self._maybe_calibrate_ensemble_temperature(X=X, y=y) y_raw = np.asarray(y, dtype=np.float64) - mean, std = np.mean(y), np.std(y) - # TODO: y_train_std_ and y_train_mean_ don't seem to be used anywhere else. - self.y_train_std_ = std.item() + 1e-20 - self.y_train_mean_ = mean.item() - y = (y - self.y_train_mean_) / self.y_train_std_ + # Shared aggregation frame; each member standardizes raw targets separately. + self.y_train_mean_ = float(np.mean(y)) + self.y_train_std_ = float(np.std(y)) + StandardizeTarget.EPSILON self._rebuild_raw_space_bardist() self._build_ensemble_preprocessor_and_executor( @@ -2008,6 +2012,8 @@ def predict_batched( # noqa: C901, PLR0912 # hold the ``target_transform`` that preprocessing fitted, which # the border mapping needs. configs=[m.config for m in members], + y_train_mean=worker.y_train_mean_, + y_train_std=worker.y_train_std_, raw_space_bardist=worker.raw_space_bardist_, znorm_space_bardist=worker.znorm_space_bardist_, X_query_raw=torch.zeros(n_test, 1), @@ -2054,6 +2060,8 @@ def predict_batched( # noqa: C901, PLR0912 config=configs[lane], znorm_borders=znorm_borders, std_borders=std_borders, + znorm_mean=group_items[lane].y_train_mean, + znorm_std=group_items[lane].y_train_std, ) previous = accumulated[lane] accumulated[lane] = ( @@ -2093,11 +2101,15 @@ def _translate_batched_logits( config: RegressorEnsembleConfig, znorm_borders: torch.Tensor, std_borders: np.ndarray, + znorm_mean: float, + znorm_std: float, ) -> torch.Tensor: """Map one estimator's output for one dataset onto the shared borders. Same border translation as :meth:`predict`, for a single - (estimator, dataset) pair of the fused forward. + (estimator, dataset) pair of the fused forward. `znorm_mean` and + `znorm_std` belong to this dataset, not to `self`: the worker is refitted + per dataset, so its own attributes only ever hold the last one's. """ out_d = output.float() temperature = resolved_softmax_temperature(self) @@ -2110,6 +2122,8 @@ def _translate_batched_logits( logit_cancel_mask, descending_borders, borders_t = transform_borders_one( std_borders, target_transform=config.target_transform, + znorm_mean=znorm_mean, + znorm_std=znorm_std, repair_nan_borders_after_transform=self.inference_config_.FIX_NAN_BORDERS_AFTER_TARGET_TRANSFORM, ) if descending_borders: @@ -2248,6 +2262,7 @@ def _iter_forward_executor( # the transformation done to the borders for a given output is dependant # upon the target_transform of the config. if config_for_ensemble.target_transform is None: + # The differentiable-input path supplies standardized targets. borders_t = std_borders.copy() logit_cancel_mask = None descending_borders = False @@ -2256,6 +2271,8 @@ def _iter_forward_executor( transform_borders_one( std_borders, target_transform=config_for_ensemble.target_transform, + znorm_mean=self.y_train_mean_, + znorm_std=self.y_train_std_, repair_nan_borders_after_transform=self.inference_config_.FIX_NAN_BORDERS_AFTER_TARGET_TRANSFORM, ) ) diff --git a/src/tabpfn/utils.py b/src/tabpfn/utils.py index b83e96ce1..d07c5b260 100644 --- a/src/tabpfn/utils.py +++ b/src/tabpfn/utils.py @@ -584,13 +584,21 @@ def transform_borders_one( borders: np.ndarray, target_transform: TransformerMixin | Pipeline, *, + znorm_mean: float, + znorm_std: float, repair_nan_borders_after_transform: bool, ) -> tuple[npt.NDArray[np.bool_] | None, bool, np.ndarray]: """Transforms the borders used for the bar distribution for regression. + The inverse returns raw target units; map them into the shared standardized + frame before applying sanity limits and aggregating predictions. + Args: borders: The borders to transform. target_transform: The target transformer to use. + znorm_mean: Mean of the training target, defining the frame the borders + are returned in. + znorm_std: Standard deviation of the training target. repair_nan_borders_after_transform: Whether to repair any borders that are NaN after the transformation. @@ -601,7 +609,12 @@ def transform_borders_one( descending_borders: Whether the borders are descending after transformation borders_t: The transformed borders themselves. """ - borders_t = target_transform.inverse_transform(borders.reshape(-1, 1)).squeeze() # type: ignore + # Use float64 for the cancelling inverse/frame operations, then restore + # the input dtype for aggregation. + borders_t = target_transform.inverse_transform( # type: ignore + borders.reshape(-1, 1).astype(np.float64) + ).squeeze() + borders_t = ((borders_t - znorm_mean) / znorm_std).astype(borders.dtype) logit_cancel_mask: npt.NDArray[np.bool_] | None = None if repair_nan_borders_after_transform: diff --git a/tests/test_batched_prediction_shape_safety.py b/tests/test_batched_prediction_shape_safety.py index 08360781d..e16c8f9c4 100644 --- a/tests/test_batched_prediction_shape_safety.py +++ b/tests/test_batched_prediction_shape_safety.py @@ -60,7 +60,12 @@ def test_classifier_heterogeneous_widths_match_serial() -> None: def test_regressor_heterogeneous_widths_match_serial() -> None: features = _features() - targets = [X[:, 5] - X[:, 6] for X in features] + targets = [ + (X[:, 5] - X[:, 6]) * scale + offset + for X, scale, offset in zip( + features, [1.0, 1000.0, 1.0], [0.0, 10000.0, 0.0], strict=True + ) + ] kwargs = { "n_estimators": 2, "device": "cpu", @@ -82,6 +87,6 @@ def test_regressor_heterogeneous_widths_match_serial() -> None: for index in range(2): reference = TabPFNRegressor(**kwargs).fit(features[index], targets[index]) np.testing.assert_allclose( - batched[index], reference.predict(features[index][:6]), atol=2e-2 + batched[index], reference.predict(features[index][:6]), rtol=2e-3, atol=2e-2 ) np.testing.assert_allclose(batched[0], batched[2], atol=2e-2) diff --git a/tests/test_downsample_correction.py b/tests/test_downsample_correction.py index b2c0892a9..e042ed653 100644 --- a/tests/test_downsample_correction.py +++ b/tests/test_downsample_correction.py @@ -21,6 +21,8 @@ temper_and_correct_logits, ) from tabpfn.finetuning.data_util import ( + ClassifierBatch, + RegressorBatch, get_preprocessed_dataset_chunks, meta_dataset_collator, ) @@ -608,6 +610,7 @@ def test__classifier__fit_from_preprocessed_clears_stale_correction(): fresh = _downsampling_classifier() batch = _finetuning_batch(stale, X, y, "classifier") + assert isinstance(batch, ClassifierBatch) cat_indices = batch.cat_indices outputs = [] for clf in (stale, fresh): @@ -630,6 +633,7 @@ def test__regressor__fit_from_preprocessed_clears_stale_correction(): fresh = _downsampling_regressor() batch = _finetuning_batch(stale, X, y, "regressor") + assert isinstance(batch, RegressorBatch) outputs = [] for reg in (stale, fresh): reg.fit_from_preprocessed( @@ -639,9 +643,12 @@ def test__regressor__fit_from_preprocessed_clears_stale_correction(): batch.configs, performance_options=PerformanceOptions(), ) - # The finetuning loop hands the batch's bar distributions to the estimator. + # The finetuning loop hands the batch's bar distributions and target frame + # to the estimator. reg.znorm_space_bardist_ = batch.znorm_space_bardist reg.raw_space_bardist_ = batch.raw_space_bardist + reg.y_train_mean_ = batch.y_train_mean + reg.y_train_std_ = batch.y_train_std assert reg.downsample_correction_log_weights_ is None out = reg.forward(batch.X_query) outputs.append((out[0] if isinstance(out, tuple) else out).detach()) diff --git a/tests/test_finetuning_regressor.py b/tests/test_finetuning_regressor.py index aca54a1fa..7770191f7 100644 --- a/tests/test_finetuning_regressor.py +++ b/tests/test_finetuning_regressor.py @@ -38,6 +38,7 @@ _targets_in_estimator_space, ) from tabpfn.preprocessing import RegressorEnsembleConfig +from tabpfn.preprocessing.target_transform import make_target_transform from tabpfn.regressor import TabPFNRegressor from tabpfn.settings import settings @@ -554,4 +555,50 @@ def test_regressor_dataset_and_collator_batches_type( assert batch.X_query_raw.shape[0] == 1 assert batch.y_query_raw.shape[0] == 1 assert batch.y_query.shape[0] == 1 + + # The frame relating the two bar distributions. `fit_from_preprocessed` + # never sees the target, so the regressor can only get it from here, and + # the forward pass needs it to map an estimator's borders back. + assert batch.y_train_std > 0.0 + np.testing.assert_allclose( + batch.raw_space_bardist.borders.cpu().numpy(), + batch.znorm_space_bardist.borders.cpu().numpy() * batch.y_train_std + + batch.y_train_mean, + rtol=1e-5, + ) break + + +@pytest.mark.parametrize( + "dtype", [torch.int32, torch.int64, torch.float32, torch.float64] +) +def test__forward_with_loss__uses_raw_targets_for_member_pipeline( + dtype: torch.dtype, +) -> None: + """The member pipeline must receive query targets in their original units.""" + y_train = torch.tensor([100, 200, 300], dtype=dtype).numpy() + pipeline = make_target_transform(None).fit(y_train.reshape(-1, 1)) + raw = torch.tensor([[150, 250]], dtype=dtype) + standardized = ((raw - y_train.mean()) / y_train.std()).float() + batch = mock.Mock( + y_query=standardized, + y_query_raw=raw, + configs=[[mock.Mock(target_transform=pipeline)]], + ) + reg = FinetunedTabPFNRegressor( + device="cpu", + n_estimators_finetune=1, + n_estimators_validation=1, + n_estimators_final_inference=1, + ) + reg._local_n_estimators_ = 1 + reg._bardist_loss = mock.Mock(num_bars=3) + logits = torch.zeros(2, 1, 3) + with ( + mock.patch.object(reg, "_training_forward", return_value=(None, [logits], [])), + mock.patch( + "tabpfn.finetuning.finetuned_regressor._compute_regression_loss" + ) as loss, + ): + reg._forward_with_loss(batch) + torch.testing.assert_close(loss.call_args.kwargs["targets_BQ"], standardized) diff --git a/tests/test_preprocessing/test_target_transform.py b/tests/test_preprocessing/test_target_transform.py new file mode 100644 index 000000000..3ffdbb84d --- /dev/null +++ b/tests/test_preprocessing/test_target_transform.py @@ -0,0 +1,151 @@ +# Copyright (c) Prior Labs GmbH 2026. + +"""Tests for the regression target pipelines.""" + +from __future__ import annotations + +import pickle + +import numpy as np +import pytest + +from tabpfn.preprocessing.steps import ( + get_all_reshape_feature_distribution_preprocessors, +) +from tabpfn.preprocessing.target_transform import ( + STANDARDIZE_STEP, + TARGET_TRANSFORM_STEP, + StandardizeTarget, + make_target_transform, +) +from tabpfn.utils import transform_borders_one + +TRANSFORM_NAMES = ["1_plus_log", "safepower", "quantile_norm", "robust", "none"] + + +def _target(n: int = 200) -> np.ndarray: + """A strictly positive, right-skewed target.""" + rng = np.random.default_rng(0) + return np.exp(rng.normal(5.0, 1.0, size=n)) + + +def _get_transform(name: str, n: int): # noqa: ANN202 + return get_all_reshape_feature_distribution_preprocessors( + num_examples=n, random_state=0 + )[name] + + +def test__standardize_target__matches_the_estimators_own_znormalisation() -> None: + """Preserve the exact arithmetic of global target standardization.""" + y = _target() + + standardized = StandardizeTarget().fit_transform(y.reshape(-1, 1)).ravel() + + expected = (y - np.mean(y)) / (np.std(y) + StandardizeTarget.EPSILON) + assert np.array_equal(standardized, expected) + + +def test__standardize_target__inverse_transform_round_trips() -> None: + y = _target() + step = StandardizeTarget().fit(y.reshape(-1, 1)) + + np.testing.assert_allclose( + step.inverse_transform(step.transform(y.reshape(-1, 1))).ravel(), y, rtol=1e-12 + ) + + +def test__standardize_target__constant_target_does_not_divide_by_zero() -> None: + """`fit` rejects a constant target, but the epsilon must hold regardless.""" + standardized = StandardizeTarget().fit_transform(np.full((10, 1), 3.0)) + + assert np.isfinite(standardized).all() + + +def test__make_target_transform__without_a_transform_only_standardizes() -> None: + y = _target() + pipeline = make_target_transform(None) + + assert list(pipeline.named_steps) == [STANDARDIZE_STEP] + expected = (y - np.mean(y)) / (np.std(y) + StandardizeTarget.EPSILON) + assert np.array_equal(pipeline.fit_transform(y.reshape(-1, 1)).ravel(), expected) + + +@pytest.mark.parametrize("name", TRANSFORM_NAMES) +def test__make_target_transform__transform_sees_the_standardized_target( + name: str, +) -> None: + """The preset receives standardized targets.""" + y = _target() + transform = _get_transform(name, len(y)) + + got = make_target_transform(transform).fit_transform(y.reshape(-1, 1)).astype(float) + + standardized = (y - np.mean(y)) / (np.std(y) + StandardizeTarget.EPSILON) + expected = _get_transform(name, len(y)).fit_transform(standardized.reshape(-1, 1)) + np.testing.assert_allclose(got, np.asarray(expected, dtype=float), rtol=1e-12) + + +@pytest.mark.parametrize("name", TRANSFORM_NAMES) +def test__make_target_transform__inverse_transform_returns_original_units( + name: str, +) -> None: + """The inverse returns raw units, independently of the aggregation frame.""" + y = _target() + pipeline = make_target_transform(_get_transform(name, len(y))) + transformed = pipeline.fit_transform(y.reshape(-1, 1)) + + np.testing.assert_allclose( + pipeline.inverse_transform(transformed).ravel(), y, rtol=1e-6 + ) + + +def test__make_target_transform__is_picklable() -> None: + """Fitted configs are pickled for joblib workers and for `save_fit_state`.""" + y = _target() + pipeline = make_target_transform(_get_transform("1_plus_log", len(y))) + pipeline.fit(y.reshape(-1, 1)) + + restored = pickle.loads(pickle.dumps(pipeline)) # noqa: S301 + + np.testing.assert_array_equal( + restored.transform(y.reshape(-1, 1)), pipeline.transform(y.reshape(-1, 1)) + ) + assert TARGET_TRANSFORM_STEP in restored.named_steps + + +def test__transform_borders_one__maps_borders_into_the_znorm_frame() -> None: + """Standardization and the aggregation-frame mapping cancel without a preset.""" + y = _target() + pipeline = make_target_transform(None) + pipeline.fit(y.reshape(-1, 1)) + borders = np.linspace(-5.0, 5.0, 101, dtype=np.float32) + + _, descending, borders_t = transform_borders_one( + borders, + target_transform=pipeline, + znorm_mean=float(np.mean(y)), + znorm_std=float(np.std(y)) + StandardizeTarget.EPSILON, + repair_nan_borders_after_transform=True, + ) + + assert not descending + np.testing.assert_allclose(borders_t, borders, atol=1e-6) + + +def test__transform_borders_one__guard_is_applied_in_the_znorm_frame() -> None: + """Border sanity limits apply in standard deviations, not raw target units.""" + y = _target() * 1e6 + pipeline = make_target_transform(None) + pipeline.fit(y.reshape(-1, 1)) + borders = np.linspace(-5.0, 5.0, 101, dtype=np.float32) + + logit_cancel_mask, _, borders_t = transform_borders_one( + borders, + target_transform=pipeline, + znorm_mean=float(np.mean(y)), + znorm_std=float(np.std(y)) + StandardizeTarget.EPSILON, + repair_nan_borders_after_transform=True, + ) + + assert logit_cancel_mask is None + np.testing.assert_allclose(borders_t, borders, atol=1e-6) diff --git a/tests/test_regressor_interface.py b/tests/test_regressor_interface.py index 0c5a0b52f..1414bd36a 100644 --- a/tests/test_regressor_interface.py +++ b/tests/test_regressor_interface.py @@ -1452,6 +1452,69 @@ def test__predict_batched__uses_fitted_target_transforms() -> None: np.testing.assert_allclose(batched[i], ref.predict(X_tests[i]), atol=1e-4) +def test__predict_batched__datasets_with_different_target_scales() -> None: + """Each dataset's borders must be mapped back in its own target frame. + + The worker is refitted per dataset, so its own `y_train_mean_`/`y_train_std_` + only ever hold the last dataset's. Targets that differ by six orders of + magnitude make a leaked frame impossible to miss. + """ + X, y = _mk_reg_dataset(0) + scales = [1e-3, 1.0, 1e6] + X_list = [X] * len(scales) + y_list = [y * scale for scale in scales] + X_tests = [X[:5]] * len(scales) + + kwargs = { + "n_estimators": 4, + "device": "cpu", + "random_state": 42, + "inference_precision": torch.float32, + } + batched = TabPFNRegressor(**kwargs).predict_batched(X_list, y_list, X_tests) + + for i, scale in enumerate(scales): + ref = TabPFNRegressor(**kwargs) + ref.fit(X_list[i], y_list[i]) + np.testing.assert_allclose( + batched[i], ref.predict(X_tests[i]), rtol=1e-4, atol=1e-6 * scale + ) + + +def test__fit__member_target_is_the_znormalized_target() -> None: + """Without subsampling or a preset, model inputs match global standardization.""" + X, y = _mk_reg_dataset(0) + reg = TabPFNRegressor( + n_estimators=1, + device="cpu", + random_state=42, + inference_config={"REGRESSION_Y_PREPROCESS_TRANSFORMS": (None,)}, + ) + reg.fit(X, y) + + member_y = np.asarray(reg.executor_.ensemble_members[0].y_train) + + assert np.array_equal(member_y, (y - np.mean(y)) / (np.std(y) + 1e-20)) + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test__fit__target_frame_is_independent_of_the_target_units(dtype: str) -> None: + """Rescaling the target changes its frame but preserves standardized inputs.""" + X, y = _mk_reg_dataset(0) + y = y.astype(dtype) + members = [] + for scale in (1.0, 1e6): + scaled_y = y * scale + reg = TabPFNRegressor(n_estimators=1, device="cpu", random_state=42) + reg.fit(X, scaled_y) + # Compute from the actual input: scaling and reduction round differently. + assert reg.y_train_mean_ == float(np.mean(scaled_y)) + assert reg.y_train_std_ == float(np.std(scaled_y)) + 1e-20 + members.append(np.asarray(reg.executor_.ensemble_members[0].y_train)) + + np.testing.assert_allclose(members[0], members[1], rtol=1e-5) + + @pytest.mark.parametrize("device", devices) def test__predict_batched__matches_per_dataset_dataframe(device: str) -> None: """Batched prediction matches per-dataset on non-numeric DataFrame inputs. From 629083b14d54be7cdcc57fa18686b97408b7d4e3 Mon Sep 17 00:00:00 2001 From: Benjamin Jaeger Date: Tue, 13 Oct 2026 12:06:59 +0200 Subject: [PATCH 2/5] fix: validate regression fit state and require target frames --- changelog/1198.changed.md | 2 +- src/tabpfn/finetuning/data_util.py | 22 +++------ src/tabpfn/finetuning/finetuned_base.py | 8 ++++ src/tabpfn/finetuning/finetuned_regressor.py | 7 +-- src/tabpfn/model_loading.py | 2 + src/tabpfn/regressor.py | 49 ++++++++++++++------ tests/test_downsample_correction.py | 33 +++++++++++-- tests/test_finetuning_regressor.py | 4 +- tests/test_save_load_fitted_model.py | 19 ++++++++ 9 files changed, 103 insertions(+), 43 deletions(-) diff --git a/changelog/1198.changed.md b/changelog/1198.changed.md index 5cb4273f1..9d0f1d1b3 100644 --- a/changelog/1198.changed.md +++ b/changelog/1198.changed.md @@ -1 +1 @@ -Regression ensembles now normalize each member’s target using its own training rows, correcting predictions when row subsampling is enabled. Without row subsampling, predictions may differ slightly due to floating-point rounding. +Regression ensembles now normalize targets using each member’s training rows; without subsampling, only floating-point differences are expected. Old fitted regressors must be refitted, and custom fine-tuning loops must pass `y_train_mean` and `y_train_std` to `fit_from_preprocessed`. diff --git a/src/tabpfn/finetuning/data_util.py b/src/tabpfn/finetuning/data_util.py index 944a8e7bc..d5e3427de 100644 --- a/src/tabpfn/finetuning/data_util.py +++ b/src/tabpfn/finetuning/data_util.py @@ -18,6 +18,7 @@ from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution from tabpfn.preprocessing.datamodel import FeatureModality, FeatureSchema from tabpfn.preprocessing.ensemble import TabPFNEnsemblePreprocessor +from tabpfn.preprocessing.target_transform import StandardizeTarget from tabpfn.utils import infer_random_state, pad_tensors if TYPE_CHECKING: @@ -81,8 +82,8 @@ class RegressorBatch: znorm_space_bardist: FullSupportBarDistribution X_query_raw: torch.Tensor y_query_raw: torch.Tensor - y_train_mean: float = 0.0 - y_train_std: float = 1.0 + y_train_mean: float + y_train_std: float Batch = TypeVar("Batch", ClassifierBatch, RegressorBatch) @@ -370,7 +371,7 @@ def __len__(self) -> int: return len(self.configs) @override - def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: C901, PLR0912 + def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: C901 """Retrieves, splits, and preprocesses the dataset config at the index. Performs train/test splitting and applies potentially multiple @@ -442,19 +443,8 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch: # noqa: # This however makes also sense when considering that # this attribute changes on every dataset if is_regression_task: - train_mean = np.mean(y_train_raw) - train_std = np.std(y_train_raw) - - eps = 1e-8 - if train_std < eps: - warnings.warn( - f"Target variable has constant or near-constant values " - f"(std={train_std:.2e}). Adding epsilon={eps} to prevent " - f"division by zero in standardization.", - UserWarning, - stacklevel=2, - ) - train_std = eps + frame = StandardizeTarget().fit(y_train_raw) + train_mean, train_std = frame.mean_, frame.std_ y_test_standardized = (y_test_raw - train_mean) / train_std znorm_borders = znorm_space_bardist_.borders diff --git a/src/tabpfn/finetuning/finetuned_base.py b/src/tabpfn/finetuning/finetuned_base.py index e4c9ce786..01e80dd8f 100644 --- a/src/tabpfn/finetuning/finetuned_base.py +++ b/src/tabpfn/finetuning/finetuned_base.py @@ -1175,6 +1175,14 @@ def _ddp_broadcast_primary_metric(metric: float) -> float: rank_batch.cat_indices, rank_batch.configs, performance_options=finetuning_performance_options, + **( + { + "y_train_mean": rank_batch.y_train_mean, + "y_train_std": rank_batch.y_train_std, + } + if isinstance(rank_batch, RegressorBatch) + else {} + ), ) if using_ddp: diff --git a/src/tabpfn/finetuning/finetuned_regressor.py b/src/tabpfn/finetuning/finetuned_regressor.py index f4f539fda..4ec4e344f 100644 --- a/src/tabpfn/finetuning/finetuned_regressor.py +++ b/src/tabpfn/finetuning/finetuned_regressor.py @@ -46,7 +46,7 @@ def _targets_in_estimator_space( transformed context targets, so their loss must use transformed targets too. """ if config.target_transform is None: - return y_Q + raise ValueError("A fitted target pipeline is required for regression loss.") y_np = config.target_transform.transform(y_Q.cpu().numpy().reshape(-1, 1)) return torch.as_tensor(y_np.ravel(), device=y_Q.device) @@ -368,13 +368,10 @@ def _should_skip_batch(self, batch: RegressorBatch) -> bool: # type: ignore[ove @override def _setup_batch(self, batch: RegressorBatch) -> None: # type: ignore[override] - """Set up bar distribution and target frame for this batch.""" + """Set up the bar distributions for this batch.""" self.finetuned_estimator_.raw_space_bardist_ = batch.raw_space_bardist self.finetuned_estimator_.bardist_ = batch.znorm_space_bardist self._bardist_loss = batch.znorm_space_bardist - # Preprocessed fitting needs the batch's target frame for border mapping. - self.finetuned_estimator_.y_train_mean_ = batch.y_train_mean - self.finetuned_estimator_.y_train_std_ = batch.y_train_std @override def _forward_with_loss(self, batch: RegressorBatch) -> torch.Tensor: # type: ignore[override] diff --git a/src/tabpfn/model_loading.py b/src/tabpfn/model_loading.py index 5952315b4..5b51dab79 100644 --- a/src/tabpfn/model_loading.py +++ b/src/tabpfn/model_loading.py @@ -1485,6 +1485,8 @@ def load_fitted_tabpfn_model( tmp / "executor_state.joblib", est.models_ ) + if saved_cls_name == "TabPFNRegressor": + est._validate_target_transform_version() est.to(device) return est diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index a02c7638f..c0c13ce79 100644 --- a/src/tabpfn/regressor.py +++ b/src/tabpfn/regressor.py @@ -35,7 +35,6 @@ from sklearn.base import ( BaseEstimator, RegressorMixin, - TransformerMixin, check_is_fitted, clone, ) @@ -1022,7 +1021,7 @@ def _initialize_dataset_preprocessing( random_state=random_state, # Use the provided rng ) # Each member fits its target pipeline on raw targets. - target_preprocessors: list[TransformerMixin | Pipeline | None] = [ + target_preprocessors: list[Pipeline] = [ make_target_transform( None if y_target_preprocessor is None @@ -1126,6 +1125,8 @@ def fit_from_preprocessed( configs: list[list[EnsembleConfig]], # Should be RegressorEnsembleConfig *, performance_options: PerformanceOptions, + y_train_mean: float, + y_train_std: float, no_refit: bool = True, ) -> TabPFNRegressor: """Used in Fine-Tuning. Fit the model to preprocessed inputs from torch @@ -1144,9 +1145,13 @@ def fit_from_preprocessed( model on each forward call inside the resulting executor. Compilation is enabled if either these options or the estimator's `enable_torch_compile` parameter enables it. + y_train_mean: Raw training-target mean for the batch. + y_train_std: Raw training-target standard deviation plus epsilon. no_refit: if True, the classifier will not be reinitialized when calling fit multiple times. """ + self.y_train_mean_ = y_train_mean + self.y_train_std_ = y_train_std if self.fit_mode != "batched": logging.warning( "The model was not in 'batched' mode. " @@ -1199,6 +1204,7 @@ def fit_from_preprocessed( ), ) + self.target_transform_version_ = 1 return self def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Self: @@ -1300,6 +1306,7 @@ def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Sel y_raw=y_float.detach().cpu().float().numpy() ) + self.target_transform_version_ = 1 return self @config_context(transform_output="default") # type: ignore @@ -1364,10 +1371,10 @@ def fit(self, X: XType, y: YType) -> Self: ) ) # No need to create an inference engine for a constant prediction + self.target_transform_version_ = 1 return self - # Must run on the raw `y`, before the z-normalisation below: each fold's - # tuning regressor derives its own mean/std from its own training split. + # Each tuning fold must derive its target statistics from its own raw targets. self._maybe_calibrate_ensemble_temperature(X=X, y=y) y_raw = np.asarray(y, dtype=np.float64) @@ -1388,6 +1395,7 @@ def fit(self, X: XType, y: YType) -> Self: ) self._compute_downsample_correction(y_raw=y_raw) + self.target_transform_version_ = 1 return self @overload @@ -1465,6 +1473,7 @@ def predict( quantiles), or a dictionary with detailed outputs. """ check_is_fitted(self) + self._validate_target_transform_version() # TODO: Move these at some point to InferenceEngine check_input_shape_matches(X, estimator=self) @@ -2046,6 +2055,8 @@ def predict_batched( # noqa: C901, PLR0912 batch.cat_indices, batch.configs, performance_options=PerformanceOptions(), + y_train_mean=batch.y_train_mean, + y_train_std=batch.y_train_std, ) accumulated: list[torch.Tensor | None] = [None] * len(group) n_estimators = 0 @@ -2116,18 +2127,18 @@ def _translate_batched_logits( if temperature != 1: out_d = out_d / temperature if config.target_transform is None: - borders_t = std_borders.copy() - logit_cancel_mask = None - else: - logit_cancel_mask, descending_borders, borders_t = transform_borders_one( - std_borders, - target_transform=config.target_transform, - znorm_mean=znorm_mean, - znorm_std=znorm_std, - repair_nan_borders_after_transform=self.inference_config_.FIX_NAN_BORDERS_AFTER_TARGET_TRANSFORM, + raise ValueError( + "A fitted target pipeline is required for batched prediction." ) - if descending_borders: - borders_t = borders_t.flip(-1) # type: ignore + logit_cancel_mask, descending_borders, borders_t = transform_borders_one( + std_borders, + target_transform=config.target_transform, + znorm_mean=znorm_mean, + znorm_std=znorm_std, + repair_nan_borders_after_transform=self.inference_config_.FIX_NAN_BORDERS_AFTER_TARGET_TRANSFORM, + ) + if descending_borders: + borders_t = borders_t.flip(-1) # type: ignore if logit_cancel_mask is not None: out_d = out_d.clone() out_d[..., logit_cancel_mask] = float("-inf") @@ -2188,6 +2199,13 @@ def _decode_batched_dataset( return logit_to_output(output_type=output_type) + def _validate_target_transform_version(self) -> None: + if getattr(self, "target_transform_version_", None) != 1: + raise ValueError( + "This fitted regressor uses an incompatible target-transform format. " + "Refit it with the current TabPFN version." + ) + def _iter_forward_executor( self, X: list[torch.Tensor] | XType, @@ -2224,6 +2242,7 @@ def _iter_forward_executor( ) check_is_fitted(self) + self._validate_target_transform_version() # Ensure torch.inference_mode is OFF to allow gradients if self.fit_mode in ["fit_preprocessors", "batched"]: # only these two modes support this option. diff --git a/tests/test_downsample_correction.py b/tests/test_downsample_correction.py index e042ed653..4e4043284 100644 --- a/tests/test_downsample_correction.py +++ b/tests/test_downsample_correction.py @@ -642,14 +642,39 @@ def test__regressor__fit_from_preprocessed_clears_stale_correction(): batch.cat_indices, batch.configs, performance_options=PerformanceOptions(), + y_train_mean=batch.y_train_mean, + y_train_std=batch.y_train_std, ) - # The finetuning loop hands the batch's bar distributions and target frame - # to the estimator. + # Custom loops provide both bar distributions before forward(). reg.znorm_space_bardist_ = batch.znorm_space_bardist reg.raw_space_bardist_ = batch.raw_space_bardist - reg.y_train_mean_ = batch.y_train_mean - reg.y_train_std_ = batch.y_train_std + assert reg.y_train_mean_ == batch.y_train_mean + assert reg.y_train_std_ == batch.y_train_std assert reg.downsample_correction_log_weights_ is None out = reg.forward(batch.X_query) outputs.append((out[0] if isinstance(out, tuple) else out).detach()) torch.testing.assert_close(outputs[0], outputs[1]) + + +@pytest.mark.parametrize("scale", [1e-10, 1e-12]) +def test__regression_batch_frame_matches_member_for_near_constant_targets( + scale: float, +) -> None: + X, y = _zero_inflated_regression(n_zeros=180, n_nonzero=20) + y = y * scale + reg = TabPFNRegressor( + n_estimators=1, + device="cpu", + random_state=0, + inference_config={"REGRESSION_Y_PREPROCESS_TRANSFORMS": (None,)}, + ) + batch = _finetuning_batch(reg, X, y, "regressor") + assert isinstance(batch, RegressorBatch) + standardize = batch.configs[0][0].target_transform.named_steps["standardize_target"] + assert batch.y_train_mean == standardize.mean_ + assert batch.y_train_std == standardize.std_ + np.testing.assert_allclose( + batch.y_query.numpy(), + (batch.y_query_raw.numpy() - standardize.mean_) / standardize.std_, + rtol=1e-6, + ) diff --git a/tests/test_finetuning_regressor.py b/tests/test_finetuning_regressor.py index 7770191f7..8371af02b 100644 --- a/tests/test_finetuning_regressor.py +++ b/tests/test_finetuning_regressor.py @@ -327,8 +327,8 @@ def test__targets_in_estimator_space__applies_the_estimator_target_transform() - """Members with a target transform are scored on transformed targets.""" y = torch.tensor([0.5, -1.0, 2.0]) - plain = _targets_in_estimator_space(y, mock.Mock(target_transform=None)) - assert plain is y + with pytest.raises(ValueError, match="fitted target pipeline"): + _targets_in_estimator_space(y, mock.Mock(target_transform=None)) negate = FunctionTransformer(lambda x: -x) transformed = _targets_in_estimator_space(y, mock.Mock(target_transform=negate)) diff --git a/tests/test_save_load_fitted_model.py b/tests/test_save_load_fitted_model.py index e311bc4bd..99bc9c670 100644 --- a/tests/test_save_load_fitted_model.py +++ b/tests/test_save_load_fitted_model.py @@ -4,6 +4,7 @@ from __future__ import annotations +import pickle import zipfile from copy import deepcopy from itertools import product @@ -608,3 +609,21 @@ def test__save_fit_state__serializes_inference_config_dataclass( loaded = estimator_class.load_from_fit_state(path, device="cpu") assert loaded.inference_config_ == model.inference_config_ _assert_roundtrip_predictions(model, loaded, X, cross_device=False) + + +@pytest.mark.parametrize("format", ["archive", "pickle"]) +def test__legacy_regressor_fit_state_requires_refitting( + tmp_path: Path, format: str +) -> None: + X, y = _make_regression_data() + model = TabPFNRegressor(n_estimators=2, device="cpu", random_state=0).fit(X, y) + del model.target_transform_version_ + if format == "archive": + path = tmp_path / "legacy.tabpfn_fit" + model.save_fit_state(path) + with pytest.raises(ValueError, match="Refit it"): + load_fitted_tabpfn_model(path, device="cpu") + else: + restored = pickle.loads(pickle.dumps(model)) # noqa: S301 + with pytest.raises(ValueError, match="Refit it"): + restored.predict(X[:4]) From 99d2660af73226e70c9a887d0abc9325169b2790 Mon Sep 17 00:00:00 2001 From: Benjamin Jaeger Date: Tue, 13 Oct 2026 12:23:45 +0200 Subject: [PATCH 3/5] docs: document fitted regressor incompatibility without version checks --- changelog/1198.breaking.md | 1 + changelog/1198.changed.md | 2 +- src/tabpfn/model_loading.py | 2 -- src/tabpfn/regressor.py | 13 ------------- tests/test_save_load_fitted_model.py | 19 ------------------- 5 files changed, 2 insertions(+), 35 deletions(-) create mode 100644 changelog/1198.breaking.md diff --git a/changelog/1198.breaking.md b/changelog/1198.breaking.md new file mode 100644 index 000000000..9c114cbf5 --- /dev/null +++ b/changelog/1198.breaking.md @@ -0,0 +1 @@ +Previously saved fitted regressors (`.tabpfn_fit` archives and pickled estimators) are incompatible with the updated target preprocessing and must be refitted. Custom fine-tuning loops must pass `y_train_mean` and `y_train_std` to `fit_from_preprocessed`. diff --git a/changelog/1198.changed.md b/changelog/1198.changed.md index 9d0f1d1b3..a47545687 100644 --- a/changelog/1198.changed.md +++ b/changelog/1198.changed.md @@ -1 +1 @@ -Regression ensembles now normalize targets using each member’s training rows; without subsampling, only floating-point differences are expected. Old fitted regressors must be refitted, and custom fine-tuning loops must pass `y_train_mean` and `y_train_std` to `fit_from_preprocessed`. +Regression ensembles now normalize targets using each member’s training rows; without subsampling, only floating-point differences are expected. diff --git a/src/tabpfn/model_loading.py b/src/tabpfn/model_loading.py index 5b51dab79..5952315b4 100644 --- a/src/tabpfn/model_loading.py +++ b/src/tabpfn/model_loading.py @@ -1485,8 +1485,6 @@ def load_fitted_tabpfn_model( tmp / "executor_state.joblib", est.models_ ) - if saved_cls_name == "TabPFNRegressor": - est._validate_target_transform_version() est.to(device) return est diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index c0c13ce79..a888b944f 100644 --- a/src/tabpfn/regressor.py +++ b/src/tabpfn/regressor.py @@ -1204,7 +1204,6 @@ def fit_from_preprocessed( ), ) - self.target_transform_version_ = 1 return self def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Self: @@ -1306,7 +1305,6 @@ def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Sel y_raw=y_float.detach().cpu().float().numpy() ) - self.target_transform_version_ = 1 return self @config_context(transform_output="default") # type: ignore @@ -1371,7 +1369,6 @@ def fit(self, X: XType, y: YType) -> Self: ) ) # No need to create an inference engine for a constant prediction - self.target_transform_version_ = 1 return self # Each tuning fold must derive its target statistics from its own raw targets. @@ -1395,7 +1392,6 @@ def fit(self, X: XType, y: YType) -> Self: ) self._compute_downsample_correction(y_raw=y_raw) - self.target_transform_version_ = 1 return self @overload @@ -1473,7 +1469,6 @@ def predict( quantiles), or a dictionary with detailed outputs. """ check_is_fitted(self) - self._validate_target_transform_version() # TODO: Move these at some point to InferenceEngine check_input_shape_matches(X, estimator=self) @@ -2199,13 +2194,6 @@ def _decode_batched_dataset( return logit_to_output(output_type=output_type) - def _validate_target_transform_version(self) -> None: - if getattr(self, "target_transform_version_", None) != 1: - raise ValueError( - "This fitted regressor uses an incompatible target-transform format. " - "Refit it with the current TabPFN version." - ) - def _iter_forward_executor( self, X: list[torch.Tensor] | XType, @@ -2242,7 +2230,6 @@ def _iter_forward_executor( ) check_is_fitted(self) - self._validate_target_transform_version() # Ensure torch.inference_mode is OFF to allow gradients if self.fit_mode in ["fit_preprocessors", "batched"]: # only these two modes support this option. diff --git a/tests/test_save_load_fitted_model.py b/tests/test_save_load_fitted_model.py index 99bc9c670..e311bc4bd 100644 --- a/tests/test_save_load_fitted_model.py +++ b/tests/test_save_load_fitted_model.py @@ -4,7 +4,6 @@ from __future__ import annotations -import pickle import zipfile from copy import deepcopy from itertools import product @@ -609,21 +608,3 @@ def test__save_fit_state__serializes_inference_config_dataclass( loaded = estimator_class.load_from_fit_state(path, device="cpu") assert loaded.inference_config_ == model.inference_config_ _assert_roundtrip_predictions(model, loaded, X, cross_device=False) - - -@pytest.mark.parametrize("format", ["archive", "pickle"]) -def test__legacy_regressor_fit_state_requires_refitting( - tmp_path: Path, format: str -) -> None: - X, y = _make_regression_data() - model = TabPFNRegressor(n_estimators=2, device="cpu", random_state=0).fit(X, y) - del model.target_transform_version_ - if format == "archive": - path = tmp_path / "legacy.tabpfn_fit" - model.save_fit_state(path) - with pytest.raises(ValueError, match="Refit it"): - load_fitted_tabpfn_model(path, device="cpu") - else: - restored = pickle.loads(pickle.dumps(model)) # noqa: S301 - with pytest.raises(ValueError, match="Refit it"): - restored.predict(X[:4]) From 55888514053f056a130452a255c0de27c3536f39 Mon Sep 17 00:00:00 2001 From: Benjamin Jaeger Date: Tue, 13 Oct 2026 12:28:34 +0200 Subject: [PATCH 4/5] test: pass regression target frames in custom loop coverage --- tests/test_fit_from_preprocessed.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/test_fit_from_preprocessed.py b/tests/test_fit_from_preprocessed.py index 28428eaa0..a50792749 100644 --- a/tests/test_fit_from_preprocessed.py +++ b/tests/test_fit_from_preprocessed.py @@ -86,17 +86,18 @@ def _forward_on_batch( estimator: TabPFNClassifier | TabPFNRegressor, batch: ClassifierBatch | RegressorBatch, ) -> list[torch.Tensor]: - """Run one step of a user-written fine-tuning loop and return its outputs. - - This is all a custom loop does before `forward`. A change that needs more - steps here breaks every such loop outside this repository. - """ + """Run one step of a custom fine-tuning loop with the batch's target frame.""" estimator.fit_from_preprocessed( batch.X_context, batch.y_context, batch.cat_indices, batch.configs, performance_options=PerformanceOptions(), + **( + {"y_train_mean": batch.y_train_mean, "y_train_std": batch.y_train_std} + if isinstance(batch, RegressorBatch) + else {} + ), ) if isinstance(estimator, TabPFNClassifier): return [estimator.forward(batch.X_query).detach()] From f48014d2cd231b63c1c97309e117c3a0a86b245f Mon Sep 17 00:00:00 2001 From: Benjamin Jaeger Date: Tue, 13 Oct 2026 12:31:03 +0200 Subject: [PATCH 5/5] test: preserve double precision in target frame assertion --- tests/test_finetuning_regressor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_finetuning_regressor.py b/tests/test_finetuning_regressor.py index 8371af02b..ee2278162 100644 --- a/tests/test_finetuning_regressor.py +++ b/tests/test_finetuning_regressor.py @@ -562,7 +562,7 @@ def test_regressor_dataset_and_collator_batches_type( assert batch.y_train_std > 0.0 np.testing.assert_allclose( batch.raw_space_bardist.borders.cpu().numpy(), - batch.znorm_space_bardist.borders.cpu().numpy() * batch.y_train_std + batch.znorm_space_bardist.borders.cpu().double().numpy() * batch.y_train_std + batch.y_train_mean, rtol=1e-5, )