Repository: PriorLabs/TabPFN

Issue:
TabPFNRegressor fails on constant input data
### Describe the bug

When fitting on constant regression target, errors

### Steps/Code to Reproduce

```
import numpy as np
from tabpfn import TabPFNRegressor

# Set random seed for reproducibility
np.random.seed(42)

# Create a synthetic dataset with constant targets
n_samples = 100
n_features = 10
constant_value = 5.0

# Generate random features
X = np.random.normal(size=(n_samples, n_features))

# Create constant target values
y = np.ones(n_samples) * constant_value

# Split data into training and testing sets
n_train = int(0.8 * n_samples)
X_train, X_test = X[:n_train], X[n_train:]
y_train, y_test = y[:n_train], y[n_train:]

print(f"Training with {n_train} samples, testing with {n_samples - n_train} samples")
print(f"Target values are constant: {constant_value}")

# Initialize and train the tabPFN regressor
model = TabPFNRegressor(device='cpu')

# Train the model
model.fit(X_train, y_train)

# Make predictions
y_pred = model.predict(X_test)

# Print a few predictions to see the results
print(f"First 5 predictions: {y_pred[:5]}")
print(f"Expected value: {constant_value}")
```

### Expected Results

No error is thrown, predict constant value

### Actual Results

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
[<ipython-input-2-68dce2b9bdbb>](https://localhost:8080/#) in <cell line: 0>()
     28 
     29 # Train the model
---> 30 model.fit(X_train, y_train)
     31 
     32 # Make predictions

3 frames
[/content/tabpfn/src/tabpfn/model/bar_distribution.py](https://localhost:8080/#) in __init__(self, borders, ignore_nan_targets)
     38         full_width = self.bucket_widths.sum()
     39 
---> 40         assert (1 - (full_width / (self.borders[-1] - self.borders[0]))).abs() < 1e-2, (
     41             f"diff: {full_width - (self.borders[-1] - self.borders[0])} with"
     42             f" {full_width} {self.borders[-1]} {self.borders[0]}"

AssertionError: diff: 0.0 with 0.0 5.0 5.0

### Versions

```shell
PyTorch version: 2.6.0+cu124
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.31.6
Libc version: glibc-2.35

Python version: 3.11.11 (main, Dec  4 2024, 08:55:07) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.5.82
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 550.54.15
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.2.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               12
On-line CPU(s) list:                  0-11
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family:                           6
Model:                                85
Thread(s) per core:                   2
Core(s) per socket:                   6
Socket(s):                            1
Stepping:                             7
BogoMIPS:                             4400.29
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat avx512_vnni md_clear arch_capabilities
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            192 KiB (6 instances)
L1i cache:                            192 KiB (6 instances)
L2 cache:                             6 MiB (6 instances)
L3 cache:                             38.5 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-11
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Vulnerable
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Vulnerable
Vulnerability Spectre v1:             Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:             Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Vulnerable; BHI: Vulnerable (Syscall hardening enabled)
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Vulnerable

Dependency Versions:
--------------------
tabpfn: 2.0.6
torch: 2.6.0+cu124
numpy: 1.26.4
scipy: 1.14.1
pandas: 2.2.2
scikit-learn: 1.5.2
typing_extensions: 4.12.2
einops: 0.8.1
huggingface-hub: 0.28.1
```


Discussion:
The root cause that I have pinpointed is that the model computes target value range by utilizing max (y) - min (y) and for constant targets, this leads to a range of zero. I would like to propose two possible approaches:

1. **Add Tiny Noise to Targets**
   This solution preserves TabPFN's learning capability while maintaining practically constant predictions. It's particularly valuable for solar forecasting applications where true constant targets are rare but minor variations exist in real-world scenarios.

2. **Skip Fitting for Constant Targets**  
   This approach provides explicit edge case handling and avoids unnecessary computation. It's especially useful when we expect truly constant adjustments (e.g., during equipment failures or maintenance periods).

The first solution offers more flexibility for real-world conditions, while the second provides computational efficiency for true constants.

**I want your opinion on the following:**  
- Which approach better aligns with our production needs?  
- Should we implement both with a configuration flag?  
- Are there other edge cases we should consider in solar forecasting?  

**Sample Implementation**  
```python
# Solution 1: Tiny noise addition
y = y_constant + np.random.normal(0, 1e-6, len(y_constant))

# Solution 2: Skip fitting
if np.ptp(y) == 0:  # If all values are identical
    return np.full_like(X_test[:, 0], y[0])
```

Would love to hear thoughts from other contributors on these approaches!




Patch:
diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py
index b6a8c94b6..9b34cd3af 100644
--- a/src/tabpfn/regressor.py
+++ b/src/tabpfn/regressor.py
@@ -18,10 +18,10 @@
 from __future__ import annotations
 
 import typing
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, Union
 from typing_extensions import Self, TypedDict, overload
 
 import numpy as np
@@ -42,8 +42,6 @@
     determine_precision,
     get_preprocessed_datasets_helper,
 )
-
-# NEW: Import the specialized inference engine
 from tabpfn.inference import InferenceEngineBatchedNoPreprocessing
 from tabpfn.model.bar_distribution import FullSupportBarDistribution
 from tabpfn.preprocessing import (
@@ -84,9 +82,24 @@
         Tags = Any
 
 
-# TypedDict definitions for prediction outputs
+# --- Prediction Output Types and Constants ---
+
+# 1. Tuples for runtime validation and internal logic.
+# These are defined directly as tuples of strings for immediate clarity.
+_OUTPUT_TYPES_BASIC = ("mean", "median", "mode")
+_OUTPUT_TYPES_QUANTILES = ("quantiles",)
+_OUTPUT_TYPES = _OUTPUT_TYPES_BASIC + _OUTPUT_TYPES_QUANTILES
+_OUTPUT_TYPES_COMPOSITE = ("full", "main")
+_USABLE_OUTPUT_TYPES = _OUTPUT_TYPES + _OUTPUT_TYPES_COMPOSITE
+
+
+# 2. Type aliases for static type checking and IDE support.
+OutputType = Literal["mean", "median", "mode", "quantiles", "full", "main"]
+"""The type hint for the `output_type` parameter in `predict`."""
+
+
 class MainOutputDict(TypedDict):
-    """Dictionary containing the main output types from the TabPFN regressor."""
+    """Specifies the return structure for `output_type="main"`."""
 
     mean: np.ndarray
     median: np.ndarray
@@ -95,12 +108,18 @@ class MainOutputDict(TypedDict):
 
 
 class FullOutputDict(MainOutputDict):
-    """Dictionary containing all outputs from the TabPFN regressor."""
+    """Specifies the return structure for `output_type="full"`."""
 
     criterion: FullSupportBarDistribution
     logits: torch.Tensor
 
 
+RegressionResultType = Union[
+    np.ndarray, list[np.ndarray], MainOutputDict, FullOutputDict
+]
+"""The type hint for the return value of the `predict` method."""
+
+
 class TabPFNRegressor(RegressorMixin, BaseEstimator):
     """TabPFNRegressor class."""
 
@@ -150,18 +169,6 @@ class TabPFNRegressor(RegressorMixin, BaseEstimator):
     preprocessor_: ColumnTransformer
     """The column transformer used to preprocess the input data to be numeric."""
 
-    # TODO: consider moving the following to constants.py
-    _OUTPUT_TYPES_BASIC = ("mean", "median", "mode")
-    """The basic output types supported by the model."""
-    _OUTPUT_TYPES_QUANTILES = ("quantiles",)
-    """The quantiles output type supported by the model."""
-    _OUTPUT_TYPES = _OUTPUT_TYPES_BASIC + _OUTPUT_TYPES_QUANTILES
-    """The output types supported by the model for the "main" output type."""
-    _OUTPUT_TYPES_COMPOSITE = ("full", "main")
-    """The composite output types supported by the model."""
-    _USABLE_OUTPUT_TYPES = _OUTPUT_TYPES + _OUTPUT_TYPES_COMPOSITE
-    """The output types supported by the model."""
-
     def __init__(  # noqa: PLR0913
         self,
         *,
@@ -183,8 +190,6 @@ def __init__(  # noqa: PLR0913
         random_state: int | np.random.RandomState | np.random.Generator | None = 0,
         n_jobs: int = -1,
         inference_config: dict | ModelInterfaceConfig | None = None,
-        # NEW: Add differentiable_input for consistency, though less used in
-        # regression FT
         differentiable_input: bool = False,
     ) -> None:
         """A TabPFN interface for regression.
@@ -419,7 +424,7 @@ def get_preprocessed_datasets(
         self,
         X_raw: XType | list[XType],
         y_raw: YType | list[YType],
-        split_fn,
+        split_fn: Callable,
         max_data_size: None | int = 10000,
     ) -> DatasetCollectionWithPreprocessing:
         """Transforms raw input data into a collection of datasets,
@@ -450,13 +455,11 @@ def get_preprocessed_datasets(
         )
 
     def _initialize_model_variables(self) -> tuple[int, np.random.Generator]:
-        """Perform initialization of the model, return determined byte_size
-        and RNG object.
-        """
+        """Initializes the model, returning byte_size and RNG object."""
         return _initialize_model_variables_helper(self, "regressor")
 
     def _initialize_dataset_preprocessing(
-        self, X: XType, y: YType, rng
+        self, X: XType, y: YType, rng: np.random.Generator
     ) -> tuple[list[RegressorEnsembleConfig], XType, YType, FullSupportBarDistribution]:
         """Prepare ensemble configs and validate X, y for one dataset/chunk.
         Handle the preprocessing of the input (X and y). We also return the
@@ -546,6 +549,7 @@ def _initialize_dataset_preprocessing(
         self.bardist_ = self.bardist_.to(self.device_)
 
         assert len(ensemble_configs) == self.n_estimators
+
         return ensemble_configs, X, y, self.bardist_
 
     def fit_from_preprocessed(
@@ -555,7 +559,7 @@ def fit_from_preprocessed(
         cat_ix: list[list[int]],
         configs: list[list[EnsembleConfig]],  # Should be RegressorEnsembleConfig
         *,
-        no_refit=True,
+        no_refit: bool = True,
     ) -> TabPFNRegressor:
         """Used in Fine-Tuning. Fit the model to preprocessed inputs from torch
         dataloader inside a training loop a Dataset provided by
@@ -577,13 +581,12 @@ def fit_from_preprocessed(
         if not hasattr(self, "model_") or not no_refit:
             byte_size, rng = self._initialize_model_variables()
         else:
-            # If model exists and no_refit is True, just get byte_size
             _, _, byte_size = determine_precision(
                 self.inference_precision, self.device_
             )
             rng = None
 
-        if not self.fit_mode == "batched":
+        if self.fit_mode != "batched":
             raise ValueError(
                 "The fit_from_preprocessed function"
                 " is only supported in the batched fit_mode."
@@ -623,6 +626,7 @@ def fit(self, X: XType, y: YType) -> Self:
         Returns:
             self
         """
+        ensemble_configs: list[RegressorEnsembleConfig]
         if not hasattr(self, "model_") or not self.differentiable_input:
             byte_size, rng = self._initialize_model_variables()
             ensemble_configs, X, y, self.bardist_ = (
@@ -636,14 +640,24 @@ def fit(self, X: XType, y: YType) -> Self:
 
         if self.fit_mode == "batched":
             raise ValueError(
-                "The fit() function is currently not supported in the batched fit_mode."
+                "The fit() function is not supported in 'batched' fit_mode."
             )
 
         assert len(ensemble_configs) == self.n_estimators
 
-        # Standardize y
-        mean = np.mean(y)
-        std = np.std(y)
+        self.is_constant_target_ = np.unique(y).size == 1
+        self.constant_value_ = y[0] if self.is_constant_target_ else None
+
+        if self.is_constant_target_:
+            self.bardist_ = FullSupportBarDistribution(
+                borders=torch.tensor(
+                    [self.constant_value_ - 1e-5, self.constant_value_ + 1e-5]
+                )
+            )
+
+            return self
+
+        mean, std = np.mean(y), np.std(y)
         self.y_train_std_ = std.item() + 1e-20
         self.y_train_mean_ = mean.item()
         y = (y - self.y_train_mean_) / self.y_train_std_
@@ -713,16 +727,9 @@ def predict(
         X: XType,
         *,
         # TODO: support "ei", "pi"
-        output_type: Literal[
-            "mean",
-            "median",
-            "mode",
-            "quantiles",
-            "full",
-            "main",
-        ] = "mean",
+        output_type: OutputType = "mean",
         quantiles: list[float] | None = None,
-    ) -> np.ndarray | list[np.ndarray] | MainOutputDict | FullOutputDict:
+    ) -> RegressionResultType:
         """Runs the forward() method and then transform the logits
         from the binning space in order to predict target variable.
 
@@ -749,14 +756,13 @@ def predict(
                 the input order.
 
         Returns:
-            The predicted target variable or a list of predictions per quantile.
+            The prediction, which can be a numpy array, a list of arrays (for
+            quantiles), or a dictionary with detailed outputs.
         """
         check_is_fitted(self)
 
         # TODO: Move these at some point to InferenceEngine
         X = validate_X_predict(X, self)
-        X = _fix_dtypes(X, cat_indices=self.inferred_categorical_indices_)
-        X = _process_text_na_dataframe(X, ord_encoder=self.preprocessor_)  # type: ignore
 
         check_is_fitted(self)
 
@@ -766,9 +772,15 @@ def predict(
             assert all(
                 (0 <= q <= 1) and (isinstance(q, float)) for q in quantiles
             ), "All quantiles must be between 0 and 1 and floats."
-        if output_type not in self._USABLE_OUTPUT_TYPES:
+        if output_type not in _USABLE_OUTPUT_TYPES:
             raise ValueError(f"Invalid output type: {output_type}")
 
+        if hasattr(self, "is_constant_target_") and self.is_constant_target_:
+            return self._handle_constant_target(X.shape[0], output_type, quantiles)
+
+        X = _fix_dtypes(X, cat_indices=self.inferred_categorical_indices_)
+        X = _process_text_na_dataframe(X, ord_encoder=self.preprocessor_)  # type: ignore
+
         # Runs over iteration engine
         (
             _,
@@ -842,7 +854,7 @@ def forward(
         X: list[torch.Tensor] | XType,
         *,
         use_inference_mode: bool = False,
-    ) -> torch.Tensor:  # FIXME: Refactor to reduce complexity,
+    ) -> tuple[torch.Tensor | None, list[torch.Tensor], list[np.ndarray]]:
         """Forward pass for TabPFNRegressor Infernce Engine.
         Used in fine-tuning and prediction. Called directly
         in FineTuning training loop or by predict() function
@@ -857,7 +869,10 @@ def forward(
             directly by user, so default should be False here.
 
         Returns:
-            Averaged logits (for fine-tuning), outputs, borders.
+            A tuple containing:
+                - Averaged logits over the ensemble (for fine-tuning).
+                - Raw outputs from each estimator in the ensemble.
+                - Borders used for each estimator.
         """
         # Scenario 1: Standard inference path
         is_standard_inference = use_inference_mode and not isinstance(
@@ -967,19 +982,37 @@ def forward(
 
         return averaged_logits, outputs, borders
 
+    def _handle_constant_target(
+        self, n_samples: int, output_type: OutputType, quantiles: list[float]
+    ) -> RegressionResultType:
+        """Handles prediction when the training target `y` was a constant value."""
+        constant_prediction = np.full(n_samples, self.constant_value_)
+        if output_type in _OUTPUT_TYPES_BASIC:
+            return constant_prediction
+        if output_type == "quantiles":
+            return [np.copy(constant_prediction) for _ in quantiles]
+
+        # Handle "main" and "full"
+        main_outputs = MainOutputDict(
+            mean=constant_prediction,
+            median=np.copy(constant_prediction),
+            mode=np.copy(constant_prediction),
+            quantiles=[np.copy(constant_prediction) for _ in quantiles],
+        )
+        if output_type == "full":
+            return FullOutputDict(
+                **main_outputs,
+                criterion=self.bardist_,
+                logits=torch.zeros((n_samples, 1)),
+            )
+        return main_outputs
+
     def get_embeddings(
         self,
         X: XType,
         data_source: Literal["train", "test"] = "test",
     ) -> np.ndarray:
-        """Get the embeddings for the input data `X`.
-
-        Parameters:
-            X (XType): The input data.
-            data_source str: Extract either the train or test embeddings
-        Returns:
-            np.ndarray: The computed embeddings for each fitted estimator.
-        """
+        """Gets the embeddings for the input data `X`."""
         return _get_embeddings(self, X, data_source)
 
 
@@ -990,17 +1023,7 @@ def _logits_to_output(
     criterion: FullSupportBarDistribution,
     quantiles: list[float],
 ) -> np.ndarray | list[np.ndarray]:
-    """Convert the logits to the specified output type.
-
-    Args:
-        output_type: The output type to convert the logits to.
-        logits: The logits to convert.
-        criterion: The criterion to use for the conversion.
-        quantiles: The quantiles to use for the conversion.
-
-    Returns:
-        The converted logits or list of converted logits.
-    """
+    """Converts raw model logits to the desired prediction format."""
     if output_type == "quantiles":
         return [criterion.icdf(logits, q).cpu().detach().numpy() for q in quantiles]
 
@@ -1016,4 +1039,4 @@ def _logits_to_output(
     else:
         raise ValueError(f"Invalid output type: {output_type}")
 
-    return output.cpu().detach().numpy()  # type: ignore
+    return output.cpu().detach().numpy()


Test patch:
diff --git a/tests/test_regressor_interface.py b/tests/test_regressor_interface.py
index b36549da3..0d99c58b1 100644
--- a/tests/test_regressor_interface.py
+++ b/tests/test_regressor_interface.py
@@ -467,3 +467,49 @@ def test_constant_feature_handling(X_y: tuple[np.ndarray, np.ndarray]) -> None:
         decimal=5,  # Allow small numerical differences
         err_msg="Predictions changed after adding constant features",
     )
+
+
+def test_constant_target(X_y: tuple[np.ndarray, np.ndarray]) -> None:
+    """Test that TabPFNRegressor predicts a constant
+    value when the target y is constant.
+    """
+    X, _ = X_y
+    y_constant = np.full(X.shape[0], 5.0)  # Create a constant target array
+
+    model = TabPFNRegressor(n_estimators=2, random_state=42)
+    model.fit(X, y_constant)
+
+    predictions = model.predict(X)
+    assert np.all(predictions == 5.0), "Predictions are not constant as expected"
+
+    # Test different output types
+    predictions_median = model.predict(X, output_type="median")
+    assert np.all(
+        predictions_median == 5.0
+    ), "Median predictions are not constant as expected"
+
+    predictions_mode = model.predict(X, output_type="mode")
+    assert np.all(
+        predictions_mode == 5.0
+    ), "Mode predictions are not constant as expected"
+
+    quantiles = model.predict(X, output_type="quantiles", quantiles=[0.1, 0.9])
+    for quantile_prediction in quantiles:
+        assert np.all(
+            quantile_prediction == 5.0
+        ), "Quantile predictions are not constant as expected"
+
+    full_output = model.predict(X, output_type="full")
+    assert np.all(
+        full_output["mean"] == 5.0
+    ), "Mean predictions are not constant as expected for full output"
+    assert np.all(
+        full_output["median"] == 5.0
+    ), "Median predictions are not constant as expected for full output"
+    assert np.all(
+        full_output["mode"] == 5.0
+    ), "Mode predictions are not constant as expected for full output"
+    for quantile_prediction in full_output["quantiles"]:
+        assert np.all(
+            quantile_prediction == 5.0
+        ), "Quantile predictions are not constant as expected for full output"