Make full regression output available in AutoTabPFNRegressor

[noahho]

Describe the bug

Describe the bug

Reported by @msalem7777 in PriorLabs/TabPFN#265

setting ignore_pretraining_limits=True on model initialization creates a model that does not have the ouput_type argument defined in the .predict method.

After re-running and digging deeper, I realized I misreported the issue. The problem isn't related to ignore_pretraining_limits=True — the error actually occurs when using AutoTabPFNRegressor, regardless of that flag. It seems the output_type argument is not supported in the ensemble version (AutoTabPFNRegressor.predict()), unlike in the base TabPFNRegressor.

Feel free to close this if that’s expected behavior — otherwise I’d be happy to revise the issue title to reflect the actual scope. Thanks again!

Steps/Code to Reproduce

reg = TabPFNRegressor(ignore_pretraining_limits=True)
reg.fit(X_train, y_train)
tabpfn_full_preds = reg.predict(X_test, output_type="full")

Expected Results

No error is thrown.

Actual Results

---

TypeError Traceback (most recent call last)
Cell In[50], line 1
----> 1 tabpfn_full_preds = reg.predict(X_test, output_type="full")

TypeError: predict() got an unexpected keyword argument 'output_type'

Versions

Collecting system and dependency information...
PyTorch version: 2.3.1+cu121
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Home (10.0.19045 64-bit)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.10 (tags/v3.9.10:f2f3f53, Jan 17 2022, 15:14:21) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1050
Nvidia driver version: 555.99
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Name: Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
Manufacturer: GenuineIntel
Family: 198
Architecture: 9
ProcessorType: 3
DeviceID: CPU0
CurrentClockSpeed: 2800
MaxClockSpeed: 2801
L2CacheSize: 1024
L2CacheSpeed: None
Revision: None

Dependency Versions:
--------------------
tabpfn: 2.0.8
torch: 2.3.1+cu121
numpy: 1.26.3
scipy: 1.13.1
pandas: 2.2.2
scikit-learn: 1.5.0
typing_extensions: 4.12.2
einops: 0.8.0
huggingface-hub: 0.29.0

Steps/Code to Reproduce

No response

Expected Results

No response

Actual Results

No response

Versions

[wangyao-1998]

Just as this issue claimed, I noticed that while the base TabPFNRegressor supports uncertainty outputs (e.g., via predict(..., output_type="quantiles", quantiles=[...]) or output_type="full"), this does not seem to work when using HPO / AutoTabPFNRegressor. The wrapper only returns point predictions, and the output_type argument is not propagated.

Could you please let me know if there are plans to add support for uncertainty outputs in the HPO / AutoTabPFN pipeline? This would be extremely valuable for applications where predictive intervals or distributional information are required.

Thank you! @noahho

[noahho]

Hey @wangyao-1998 @msalem7777

Thank you for following up on this.

The .fit() method of TunedTabPFNRegressor stores the best-found TabPFNRegressor instance in the .best_model_ attribute. You can then call .predict() on this underlying model, which fully supports the output_type argument.

Here is a complete, runnable script demonstrating this approach with the improvements you suggested. This script first runs the hyperparameter search and then uses the best-found model to generate full distributional predictions, including quantiles and standard deviations.

import numpy as np
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split

from tabpfn_extensions.hpo import TunedTabPFNRegressor
from tabpfn_extensions.hpo.search_space import get_param_grid_hyperopt

# --- 1. Load and prepare regression data ---
# We use the California Housing dataset as a proper regression example.
X, y = fetch_california_housing(return_X_y=True)
X, y = X[:1000], y[:1000] # reduce data size

# Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.25,
    random_state=42,
)

# --- 2. Set up the Hyperparameter Optimization ---
# Get the default hyperopt search space for a regression task.
search_space = get_param_grid_hyperopt("regression")
# We remove the dt_pfn model from search space which doesnt allow the full prediction quantiles
del search_space['model_type']

# Initialize the TunedTabPFNRegressor.
# This will automatically run HPO during the .fit() call.
reg = TunedTabPFNRegressor(
    n_trials=5,  # Number of optimization rounds (kept low for a quick demo)
    random_state=42,
    verbose=True,
    search_space=search_space
)

# Fit the model to find the best hyperparameters.
reg.fit(X_train, y_train)

# --- 3. Access the Best Model for Full Predictions ---
# The key to the workaround is to use the .best_model_ attribute.
# This attribute holds the best-performing TabPFNRegressor instance found during the search.
best_underlying_model = reg.best_model_

# Now, call .predict() on this underlying model, which supports the `output_type` argument.
# This allows you to get quantiles, standard deviations, etc.
predictions_full = best_underlying_model.predict(X_test, output_type="full")

[wangyao-1998]

Thank you very much for the quick response and for sharing the sample code — it worked on my side.