Add CP_MDA-exact algorithm for correct uncertainty estimates with missing data.

[Fdvanleeuwen]

The CP_MDA-exact algorithm [1] allows users to obtain conformalised uncertainty estimates for TabPFNRegressor. This allows users to have correct uncertainty estimates based on the missing masks, as described in [2]. There are two functions:

CP_MDA_TabPFNRegressor: for fitting TabPFNRegressor and obtaining the correction terms for each mask.
CP_MDA_TabPFNRegressor_newdata: for applying the fitted model (and correction terms) to new test cases.

[1] Margaux Zaffran, Aymeric Dieuleveut, Julie Josse, and Yaniv Romano. Conformal prediction with missing
values. In International Conference on Machine Learning, pages 40578–40604. PMLR, 2023.
[2] van Leeuwen, F. D. Conformal Prediction for Tabular Prior-Data Fitted Networks with Missing data. In EurIPS 2025 Workshop: AI for Tabular Data.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[CLAassistant]

All committers have signed the CLA.

[gemini-code-assist]

Summary of Changes

Hello @Fdvanleeuwen, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates the CP_MDA-exact algorithm into the tabpfn_extensions library, enabling TabPFNRegressor to generate accurate and calibrated uncertainty estimates even in the presence of missing data. The solution provides a robust framework for handling various missing data patterns by calculating specific correction terms, thereby improving the reliability of predictions for tabular data.

Highlights

• New Algorithm for Missing Data: Introduces the CP_MDA-exact algorithm to provide conformalized uncertainty estimates for TabPFNRegressor when data is missing.
• Two-Step Process: Implements two core functions: CP_MDA_TabPFNRegressor for model fitting and correction term calculation, and CP_MDA_TabPFNRegressor_newdata for applying these corrections to new datasets.
• Research-Backed: The algorithm is based on recent research, specifically referencing papers by Zaffran et al. (2023) and van Leeuwen (EurIPS 2025).

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces the CP_MDA-exact algorithm for handling missing data with TabPFN, which is a valuable addition. The implementation is well-structured into two classes for training and prediction. I've found a critical issue in the data selection logic that could lead to incorrect calibration, which needs to be addressed. Additionally, I've provided feedback on improving docstrings, code style consistency, and the flexibility of the implementation. The tests are a good start, but one of them could be made more robust. The example script is clear and helpful.

[gemini-code-assist]

The TabPFNRegressor is instantiated with its default parameters, which limits flexibility. Consider allowing users to pass custom arguments to TabPFNRegressor through the CP_MDA_TabPFNRegressor constructor.

You could modify the __init__ method to accept **tabpfn_kwargs and store them:

def __init__(self, X_train, Y_train, quantiles, val_size, seed, **tabpfn_kwargs):
    # ...
    self.tabpfn_kwargs = tabpfn_kwargs

Then, you can use these arguments when creating the instance here.

Suggested change

	m_fit = TabPFNRegressor()
	m_fit = TabPFNRegressor(**getattr(self, 'tabpfn_kwargs', {}))

[bejaeger]

Hi @Fdvanleeuwen ,
thank you so much for this contribution. It looks great all in all. There are a few stylistic things I commented on that would be great if you could change them. The CI also shows some errors. If you could fix these I can have a closer look at the method.
Thanks again!

[Fdvanleeuwen]

Hi @Fdvanleeuwen , thank you so much for this contribution. It looks great all in all. There are a few stylistic things I commented on that would be great if you could change them. The CI also shows some errors. If you could fix these I can have a closer look at the method. Thanks again!

Thanks @bejaeger for the detailed review. I'm not super familiar with python conventions, hope it is a bit better now :)

[Fdvanleeuwen]

Hi @Fdvanleeuwen , thank you so much for this contribution. It looks great all in all. There are a few stylistic things I commented on that would be great if you could change them. The CI also shows some errors. If you could fix these I can have a closer look at the method. Thanks again!

Thanks @bejaeger for the detailed review. I'm not super familiar with python conventions, hope it is a bit better now :)

I worked on the formatting problems, and I guess there were issues with changing all files to lower case. All tests work locally, so hopefully it will be better this time. @bejaeger

[Fdvanleeuwen]

Hi @Fdvanleeuwen , thank you so much for this contribution. It looks great all in all. There are a few stylistic things I commented on that would be great if you could change them. The CI also shows some errors. If you could fix these I can have a closer look at the method. Thanks again!

Thanks @bejaeger for the detailed review. I'm not super familiar with python conventions, hope it is a bit better now :)

I worked on the formatting problems, and I guess there were issues with changing all files to lower case. All tests work locally, so hopefully it will be better this time. @bejaeger

Again some small formatting issues. But main problem is that many tests still fail. It might be because TabPFN v2.5 cannot be imported. Any idea how I can resolve this? @bejaeger