TabPFN Extensions

[!WARNING]

Experimental Code Notice

Please note that the extensions in this repository are experimental.

• They are less rigorously tested than the core tabpfn library.
• APIs are subject to change without notice in future releases. We welcome your feedback and contributions to help improve and stabilize them!

Interactive Notebook Tutorial

[!TIP]

Dive right in with our interactive Colab notebook! It's the best way to get a hands-on feel for TabPFN, walking you through installation, classification, and regression examples.

Installation

# Clone and install the repository
pip install "tabpfn-extensions[all] @ git+https://github.com/PriorLabs/tabpfn-extensions.git"

Available Extensions

• interpretability: Explain TabPFN predictions with SHAP values and feature selection
• many_class: Handle classification problems with more classes than your TabPFN checkpoint supports
• unsupervised: Data generation and outlier detection
• embedding: Get TabPFN's internal dense sample embeddings
• tabebm: Data augmentation using TabPFN-based Energy-Based Models
• pval_crt: Statistical feature relevance testing (p-values)
• bayesian_optimization: Bayesian optimization with TabPFN as the surrogate model and differentiable Expected Improvement

See the Documentation section below for guides, examples, and per-extension READMEs.

Backend Options

Many TabPFN Extensions works with two TabPFN implementations:

1. TabPFN Package — Full PyTorch implementation for local inference:

   pip install tabpfn
   

2. TabPFN Client — Lightweight API client for cloud-based inference:

   pip install "tabpfn-client>=0.2.7"
   

Choose the backend that fits your needs - most extensions work with either option!

Exceptions to this are embedding and bayesian_optimization, which only work with the local tabpfn package.

Documentation

Documentation for tabpfn-extensions is spread across several sources. If you are new to the project, the examples are usually the fastest way to get started; for deeper conceptual guides, see the TabPFN Docs pages.

Examples

Runnable scripts and notebooks for extensions and general use cases live in the examples/ directory of this repository:

• bayesian_optimization/ — Bayesian optimization with TabPFN as the surrogate and differentiable Expected Improvement
• embedding/ — access TabPFN's internal dense sample embeddings
• interpretability/ — SHAP values, partial dependence plots, feature selection
• many_class/ — classification with more classes than your checkpoint supports
• predictive_distribution/ — visualize the full predictive distribution from TabPFNRegressor and derive point estimates / credible intervals
• pval_crt/ — statistical feature relevance testing
• survival/ — survival analysis
• tabebm/ — data augmentation via TabEBM
• unsupervised/ — data generation, imputation, and outlier detection

TabPFN Docs pages

In-depth guides for selected extensions are available on docs.priorlabs.ai:

• Many-class

Per-extension READMEs

Some extensions ship a dedicated README alongside their source code:

• interpretability/
• pval_crt/
• tabebm/

Interactive notebook

The main TabPFN demo notebook also covers several extensions — in particular the unsupervised and interpretability extensions:

License

This project is licensed under the Apache License 2.0 — see the LICENSE file for details.

Telemetry

For details on telemetry, please see our Telemetry Reference and our Privacy Policy.

For Contributors

Interested in adding your own extension? We welcome contributions!

We use uv to manage the project's environment, so install that first.

# Clone and set up for development
git clone https://github.com/PriorLabs/tabpfn-extensions.git
cd tabpfn-extensions
uv sync
source .venv/bin/activate

# If you add optional dependencies for your extension in pyproject.toml, install them
# like this
uv sync --extra [your extension name]

# Test your extension with fast mode
FAST_TEST_MODE=1 pytest tests/test_your_extension.py -v

See our Contribution Guide for more details.

Built with ❤️ by the TabPFN community