shapiq.imputer.TabPFNImputer

class shapiq.imputer.TabPFNImputer(model, x_train, y_train, *, x_test=None, empty_prediction=None, verbose=False, predict_function=None)

Bases: Imputer[SklearnLikeModel]

An Imputer for TabPFN using the Remove-and-Contextualize paradigm.

The remove-and-contextualize paradigm is a strategy to explain the predictions of a TabPFN [2] model which uses in-context learning for prediction. Instead of imputing missing features, the TabPFNImputer removes feature columns missing in a coalition from training data and re-“trains” re-contextualizes the model with the remaining features. The model is then used to predict the data point which is also missing the features. This pardigm is described in Rundel et al. (2024) [1].

Variables:

• x_train – The training data to contextualize the model on.

• y_train – The training labels to contextualize the model on.

• empty_prediction – The model’s average prediction on an empty data point.

Parameters:

• model (TModel)

• x_train (np.ndarray)

• y_train (np.ndarray)

• x_test (np.ndarray | None)

• empty_prediction (float)

• verbose (bool)

• predict_function (Callable[[TabPFNClassifier | TabPFNRegressor, np.ndarray], np.ndarray] | None)

References

[1]

Rundel, D., Kobialka, J., von Crailsheim, C., Feurer, M., Nagler, T., Rügamer, D. (2024). Interpretable Machine Learning for TabPFN. In: Longo, L., Lapuschkin, S., Seifert, C. (eds) Explainable Artificial Intelligence. xAI 2024. Communications in Computer and Information Science, vol 2154. Springer, Cham. https://doi.org/10.1007/978-3-031-63797-1_23

[2]

Hollmann, N., Müller, S., Purucker, L. et al. Accurate predictions on small data with a tabular foundation model. Nature 637, 319-326 (2025). https://doi.org/10.1038/s41586-024-08328-6

An Imputer for TabPFN using the Remove-and-Contextualize paradigm.

Parameters:

• model (TabPFNClassifier | TabPFNRegressor) – The model to be explained as a callable function expecting data points as input and returning 1-dimensional predictions.

• x_train (ndarray) – The training data to “train” the model on. Note that the model is not actually trained but the correct train data with the correct features per coalition are put into TabPFN’s context.

• y_train (ndarray) – The training labels to “train” the model on. Note that the model is not trained but the correct train data and labels are put into TabPFN’s context.

• x_test (ndarray | None) – The test data to evaluate the model’s average (empty) prediction on. If no test data is provided, the empty prediction must be given. Defaults to None.

• empty_prediction (float | None) – The model’s average prediction on an empty data point (all features missing). This can be computed by averaging the model’s predictions on the test data.

• verbose (bool) – A flag to enable verbose output. Defaults to False.

• predict_function (Callable[[TabPFNClassifier | TabPFNRegressor, ndarray], ndarray] | None) – A function to use for prediction. If the model is not instantiated via a shapiq.Explainer object, this function must be provided. The function must accept the model and the data point as input and return the model’s predictions. If the model is instantiated via a shapiq.Explainer object, this function is automatically set to the model’s prediction function. Defaults to None.

value_function(coalitions)

The value function performs the remove-and-contextualize strategy for TabPFN.

The value function removes absent features from a coalition by “training” the model again on the subset of features. The model is then used to predict the data point with the missing features.

Parameters:

coalitions (ndarray) – A boolean array indicating which features are present (True) and which are missing (False). The shape of the array must be (n_subsets, n_players).

Return type:

ndarray

Returns:

The model’s predictions on the restricted data points. The shape of the array is

(n_subsets,).