Utilities

Installation

pip install tabpfn-extensions

`get_tags`

View source

Get estimator tags in a consistent format across different sklearn versions. This function provides compatibility between sklearn versions before and after 1.6. It returns either a Tags object (sklearn >= 1.6) or a converted Tags object from the dictionary format (sklearn < 1.6) containing metadata about the estimator’s requirements and capabilities.

get_tags(
    estimator,
)

Parameters

Parameter	Type	Default	Description
`estimator`	`estimator object`	Required	A scikit-learn estimator instance.

Returns

Type	Description
`Tags`	`tags`: An object containing metadata about the estimator’s requirements and capabilities (e.g., input types, fitting requirements, classifier/regressor specific tags).

`ClientTabPFNClassifier`

View source

`ClientTabPFNClassifier.get_params`

View source

Return parameters for this estimator.

ClientTabPFNClassifier.get_params(
    deep: bool = True,
) -> dict[str, Any]

Parameters

Parameter	Type	Default	Description
`deep`	`bool`	`True`	—

Returns

Type	Description
`dict[str, Any]`	—

`ClientTabPFNRegressor`

View source

`ClientTabPFNRegressor.get_params`

View source

Return parameters for this estimator.

ClientTabPFNRegressor.get_params(
    deep: bool = True,
) -> dict[str, Any]

Parameters

Parameter	Type	Default	Description
`deep`	`bool`	`True`	—

Returns

Type	Description
`dict[str, Any]`	—

`FakeTorchDevice`

View source

Fake used to represent torch.device used when PyTorch is not installed.

FakeTorchDevice(
    type: str,
)

Parameters

Parameter	Type	Default	Description
`type`	`str`	Required	—

`TabPFNEstimator`

View source

`TabPFNEstimator.fit`

View source

`TabPFNEstimator.predict`

View source

`get_max_num_classes`

View source

Infer the max number of classes a TabPFN estimator can predict in one fit. This is the single source of truth for the TabPFN class-count limit across tabpfn-extensions, so a fix here propagates everywhere (unsupervised classifier/regressor routing, the many-class output-coding wrapper, …). The value is read from the model’s inference config (get_inference_config().MAX_NUMBER_OF_CLASSES), which TabPFN exposes as of v8.0.0 (the minimum version this package depends on).

get_max_num_classes(
    model: Any,
) -> int | None

Parameters

Parameter	Type	Default	Description
`model`	`Any`	Required	A (typically TabPFN) classifier instance.

Returns

Type	Description
`int \| None`	The maximum number of classes the model supports, or `None` if `model` is not a TabPFN estimator (i.e. has no inherent class limit).

`get_tabpfn_models`

View source

Get the TabPFN model classes for the selected backend. USE_TABPFN_LOCAL selects the backend; the function does not silently fall back to the other one:

USE_TABPFN_LOCAL is True -> the standard tabpfn package
USE_TABPFN_LOCAL is False -> the tabpfn-client API backend

If the selected backend is not installed, an ImportError is raised naming that backend, rather than quietly using the other one.

get_tabpfn_models() -> tuple[type, type]

Returns

Type	Description
`tuple[type, type]`	tuple[type, type]: A tuple containing (TabPFNClassifier, TabPFNRegressor) classes

Raises ImportError If the selected TabPFN backend is not installed

`infer_categorical_features`

View source

Infer which columns are categorical features (constraint (a) only). This answers a data question — “is this column categorical?” — and is deliberately independent of any model constraint. Whether a categorical column has few enough levels for a TabPFN classifier to predict it (constraint (b)) is a separate concern; derive that limit with get_max_num_classes and apply it at the point of use. A column is treated as categorical if any of these hold:

It is in the caller-provided categorical_features list.
It has a string/object/category dtype (pandas DataFrame).
It contains string values (numpy object array).
It is low-cardinality: at most MAX_UNIQUE_VALUES_FOR_CATEGORICAL unique values, with more than MIN_SAMPLES_PER_CATEGORY samples per unique value on average, to avoid mislabelling columns that only look low-cardinality because the sample is too thin per level.

infer_categorical_features(
    X: np.ndarray,
    categorical_features: list[int] | None = None,
) -> list[int]

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray`	Required	The input data.
`categorical_features`	`list[int] \| None`	`None`	Initial list of categorical feature indices. If `None`, will start with an empty list.

Returns

Type	Description
`list[int]`	list[int]: The indices of the categorical features.

`infer_device`

View source

Where TabPFN itself runs: a CPU stand-in when the client serves the model.

infer_device(
    device: DeviceSpecification,
) -> torch.device | FakeTorchDevice

Type aliases

DeviceSpecification = Literal["auto", "cuda", "cpu"]

Parameters

Parameter	Type	Default	Description
`device`	`DeviceSpecification`	Required	—

Returns

Type	Description
`torch.device \| FakeTorchDevice`	—

`infer_torch_device`

View source

Where torch work of this process runs, with or without the local tabpfn. With tabpfn installed this is TabPFN’s own reading of device. Without it, the same rule on torch directly: for "auto", CUDA, else MPS, else the CPU, minus what TABPFN_EXCLUDE_DEVICES names; anything else is parsed as a torch device, the first of several.

infer_torch_device(
    device: Any,
) -> torch.device

Parameters

Parameter	Type	Default	Description
`device`	`Any`	Required	—

Returns

Type	Description
`torch.device`	—

`is_tabpfn`

View source

Check if an estimator is a TabPFN model.

is_tabpfn(
    estimator: Any,
) -> bool

Parameters

Parameter	Type	Default	Description
`estimator`	`Any`	Required	—

Returns

Type	Description
`bool`	—

`product_dict`

View source

Cartesian product of a dictionary of lists. This function takes a dictionary where each value is a list, and returns an iterator over dictionaries where each key is mapped to one element from the corresponding list.

product_dict(
    d: dict[str, list[T]],
) -> Iterator[dict[str, T]]

Parameters

Parameter	Type	Default	Description
`d`	`dict[str, list[T]]`	Required	A dictionary mapping keys to lists of values.

Returns

Type	Description
`Iterator[dict[str, T]]`	An iterator over dictionaries, each being one element of the cartesian product of the input dictionary.

Example

>>> list(product_dict({'a': [1, 2], 'b': ['x', 'y']}))
[{'a': 1, 'b': 'x'}, {'a': 1, 'b': 'y'}, {'a': 2, 'b': 'x'}, {'a': 2, 'b': 'y'}]

`softmax`

View source

Apply softmax function to convert logits to probabilities.

softmax(
    logits: NDArray,
) -> NDArray

Parameters

Parameter	Type	Default	Description
`logits`	`NDArray`	Required	Input logits array of shape (n_samples, n_classes) or (n_classes,)

Returns

Type	Description
`NDArray`	Probabilities where values sum to 1 across the last dimension

`warn_if_no_kv_cache`

View source

Warn if a TabPFN model isn’t configured to use the KV cache. The KV cache (improved with TabPFN-3) caches the encoder pass over the training set so that repeated predicts against the same fitted model don’t re-encode the training data each time. Extensions that issue many predicts per fit (e.g. imputation-based SHAP, certain feature-selection or HPO routines) benefit from it — without the cache, the encoder pass over the training set runs on every predict and these extensions can be 10-100x slower than necessary. What has to hold depends on the backend. An endpoint-backed estimator (self-hosted container, SageMaker, Foundry) needs use_kv_cache=True, which is the only condition. A local model needs both:

model was constructed with fit_mode="fit_with_cache" (a constructor argument, must be set BEFORE .fit()).
model.executor_.keep_cache_on_device is True (set AFTER .fit(); usually the default but worth setting explicitly).

This helper warns if either is missing, but does not raise — users may have intentional reasons (e.g. memory).

warn_if_no_kv_cache(
    model: Any,
    *,
    context: str = "This operation",
) -> None

Parameters

Parameter	Type	Default	Description
`model`	`Any`	Required	The TabPFN model (classifier or regressor) to inspect.
`context`	`str`	`"This operation"`	Short noun phrase describing the caller’s operation, used to make the warning message specific (e.g. `"Imputation-based SHAP"`, `"Sequential feature selection"`). Defaults to a generic `"This operation"`.

Returns

Type	Description
`None`	—

Reference

tabpfn

tabpfn-client

tabpfn-extensions

REST API

Installation

`get_tags`

`ClientTabPFNClassifier`

`ClientTabPFNClassifier.get_params`

`ClientTabPFNRegressor`

`ClientTabPFNRegressor.get_params`

`FakeTorchDevice`

`TabPFNEstimator`

`TabPFNEstimator.fit`

`TabPFNEstimator.predict`

`get_max_num_classes`

`get_tabpfn_models`

`infer_categorical_features`

`infer_device`

`infer_torch_device`

`is_tabpfn`

`product_dict`

`softmax`

`warn_if_no_kv_cache`

​Installation

​get_tags

​ClientTabPFNClassifier

​ClientTabPFNClassifier.get_params

​ClientTabPFNRegressor

​ClientTabPFNRegressor.get_params

​FakeTorchDevice

​TabPFNEstimator

​TabPFNEstimator.fit

​TabPFNEstimator.predict

​get_max_num_classes

​get_tabpfn_models

​infer_categorical_features

​infer_device

​infer_torch_device

​is_tabpfn

​product_dict

​softmax

​warn_if_no_kv_cache

Installation

`get_tags`

`ClientTabPFNClassifier`

`ClientTabPFNClassifier.get_params`

`ClientTabPFNRegressor`

`ClientTabPFNRegressor.get_params`

`FakeTorchDevice`

`TabPFNEstimator`

`TabPFNEstimator.fit`

`TabPFNEstimator.predict`

`get_max_num_classes`

`get_tabpfn_models`

`infer_categorical_features`

`infer_device`

`infer_torch_device`

`is_tabpfn`

`product_dict`

`softmax`

`warn_if_no_kv_cache`