Imputation and data generation

Looking for usage documentation? Check out Data generation.

`TabPFNUnsupervisedModel.generate_synthetic_data`

Generate synthetic tabular data samples using the fitted TabPFN models. This method uses imputation to create synthetic data, starting with a matrix of NaN values and filling in each feature sequentially. Samples are generated feature by feature in a single pass, with each feature conditioned on previously generated features.

TabPFNUnsupervisedModel.generate_synthetic_data(
    n_samples: int = 100,
    t: float = 1.0,
    n_permutations: int = 3,
    dag: dict[int, list[int]] | None = None,
) -> torch.Tensor

Parameters

Parameter	Type	Default	Description
`n_samples`	`int`	`100`	int, default=100 Number of synthetic samples to generate
`t`	`float`	`1.0`	float, default=1.0 Temperature parameter for sampling. Controls randomness: - Higher values (e.g., 1.0) produce more diverse samples - Lower values (e.g., 0.1) produce more deterministic samples
`n_permutations`	`int`	`3`	int, default=3 Number of feature permutations to use for generation More permutations may provide more robust results but increase computation time
`dag`	`dict[int, list[int]] \| None`	`None`	dict[int, list[int]] \| `None`, default=`None` Optional Directed Acyclic Graph mapping each column index to the list of column indices it depends on. When provided, columns are generated in topological order and each column is conditioned on its DAG parents only. Every column must appear as a key; map a column to an empty list to sample it marginally. Useful for causally-informed synthesis.

Returns

Type	Description
`torch.Tensor`	torch.Tensor: Generated synthetic data of shape (`n_samples`, n_features)

Raises AssertionError If the model is not fitted (self.X_ does not exist) ValueError If dag contains a cycle or does not specify every feature

`TabPFNUnsupervisedModel.impute`

View source

Impute missing values in the input data using the fitted TabPFN models. This method fills missing values (np.nan) in the input data by predicting each missing value based on the observed values in the same sample. The imputation uses multiple random feature permutations to improve robustness.

TabPFNUnsupervisedModel.impute(
    X: torch.Tensor | np.ndarray | pd.DataFrame,
    t: float = 1e-09,
    n_permutations: int = 10,
    dag: dict[int, list[int]] | None = None,
) -> torch.Tensor

Parameters

Parameter	Type	Default	Description
`X`	`torch.Tensor \| np.ndarray \| pd.DataFrame`	Required	Union[torch.Tensor, `np.ndarray`, `pd.DataFrame`] Input data of shape (n_samples, n_features) with missing values encoded as np.nan.
`t`	`float`	`1e-09`	float, default=0.000000001 Temperature for sampling from the imputation distribution. Lower values result in more deterministic imputations, while higher values introduce more randomness.
`n_permutations`	`int`	`10`	int, default=10 Number of random feature permutations to use for imputation. Higher values may improve robustness but increase computation time.
`dag`	`dict[int, list[int]] \| None`	`None`	dict[int, list[int]] \| `None`, default=`None` Optional Directed Acyclic Graph mapping each column index to the list of column indices it depends on. When provided, columns are imputed in topological order and each column is conditioned on its DAG parents instead of all other features. Every column must appear as a key; map a column to an empty list to impute it without conditioning. Useful for causally-informed imputation.

Returns

Type	Description
`torch.Tensor`	torch.Tensor Imputed data with missing values replaced, of shape (n_samples, n_features).

Note The model must be fitted with training data before calling this method.

`TabPFNUnsupervisedModel.impute_`

View source

Impute missing values (np.nan) in X by sampling all cells independently from the trained models.

TabPFNUnsupervisedModel.impute_(
    X: torch.Tensor,
    t: float = 1e-09,
    n_permutations: int = 10,
    condition_on_all_features: bool = True,
    dag: dict[int, list[int]] | None = None,
    fast_mode: bool = False,
) -> torch.Tensor

Parameters

Parameter	Type	Default	Description
`X`	`torch.Tensor`	Required	torch.Tensor Input data of shape (n_samples, n_features) with missing values encoded as np.nan
`t`	`float`	`1e-09`	float, default=0.000000001 Temperature for sampling from the imputation distribution, lower values are more deterministic
`n_permutations`	`int`	`10`	int, default=10 Number of permutations to use for imputation
`condition_on_all_features`	`bool`	`True`	bool, default=`True` Whether to condition on all other features (`True`) or only previous features (`False`)
`dag`	`dict[int, list[int]] \| None`	`None`	dict[int, list[int]] \| `None`, default=`None` Optional Directed Acyclic Graph mapping each column index to its list of parent column indices (i.e. the features it depends on). When provided, columns are imputed in topological order and each column is conditioned on exactly its DAG parents. Mutually exclusive with `condition_on_all_features=True`. Every feature must appear as a key; map a feature to an empty list to impute it without conditioning.
`fast_mode`	`bool`	`False`	bool, default=`False` Whether to use faster settings for testing

Returns

Type	Description
`torch.Tensor`	torch.Tensor: Imputed data with missing values replaced

Raises ValueError If dag is combined with condition_on_all_features=True, contains a cycle, or does not specify every feature.

`TabPFNUnsupervisedModel.impute_single_permutation_`

View source

Impute missing values (np.nan) in X by sampling all cells independently from the trained models.

TabPFNUnsupervisedModel.impute_single_permutation_(
    X: torch.Tensor,
    feature_permutation: list[int] | tuple[int, ...],
    t: float = 1e-09,
    condition_on_all_features: bool = True,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]

Parameters

Parameter	Type	Default	Description
`X`	`torch.Tensor`	Required	Input data of the shape (num_examples, num_features) with missing values encoded as np.nan
`feature_permutation`	`list[int] \| tuple[int, …]`	Required	—
`t`	`float`	`1e-09`	Temperature for sampling from the imputation distribution, lower values are more deterministic
`condition_on_all_features`	`bool`	`True`	—

Returns

Type	Description
`tuple[torch.Tensor, dict[str, torch.Tensor]]`	Imputed data, with missing values replaced

`TabPFNUnsupervisedModel.sample_from_model_prediction_`

View source

Sample values from a model’s prediction distribution.

TabPFNUnsupervisedModel.sample_from_model_prediction_(
    column_idx: int,
    X_fit: torch.Tensor,
    model: Any,
    X_predict: torch.Tensor,
    t: float,
) -> tuple[dict[str, Any] | torch.Tensor, torch.Tensor]

Parameters

Parameter	Type	Default	Description
`column_idx`	`int`	Required	Index of the column being predicted
`X_fit`	`torch.Tensor`	Required	Training data used to determine feature type
`model`	`Any`	Required	The trained model (classifier or regressor)
`X_predict`	`torch.Tensor`	Required	Input data for prediction
`t`	`float`	Required	Temperature parameter for sampling (lower values = more deterministic)

Returns

Type	Description
`tuple[dict[str, Any] \| torch.Tensor, torch.Tensor]`	tuple containing: - The raw prediction output (dictionary for regressors, tensor for classifiers) - The sampled values as a tensor

`impute_column`

View source

Impute column col of X (np.nan = missing) in place. Uses every other column as features (feature NaNs left in — TabPFN handles them), fits model on the rows where col is observed, and predicts the rows where it is missing.

impute_column(
    X: np.ndarray,
    col: int,
    model: object,
    *,
    classification: bool = False,
) -> np.ndarray

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray`	Required	Data of shape `(n_samples, n_features)` with missing values as `np.nan`.
`col`	`int`	Required	Index of the column to impute.
`model`	`object`	Required	A fitted-on-call TabPFN estimator (regressor, or classifier when `classification=True`).
`classification`	`bool`	`False`	If `True`, treat the target as categorical (fit on integer labels and predict class labels).

Returns

Type	Description
`np.ndarray`	`np.ndarray`: `X` with column `col` filled (mutated in place and returned).

`simple_impute`

View source

Impute all missing values (np.nan) in X, one column at a time. For each column that contains missing values, fit a TabPFN model on the rows where that column is observed and predict the missing rows, using all other columns as features. Numerical columns use a regressor; categorical columns use a classifier.

simple_impute(
    X: np.ndarray,
    tabpfn_clf: object | None = None,
    tabpfn_reg: object | None = None,
    categorical_features: list[int] | None = None,
    use_imputed: bool = False,
) -> np.ndarray

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray`	Required	Data of shape `(n_samples, n_features)` with missing values encoded as `np.nan`. Converted to a float array (a copy is made; the input is not modified).
`tabpfn_clf`	`object \| None`	`None`	TabPFN classifier used for categorical columns. Defaults to `TabPFNClassifier()` when needed.
`tabpfn_reg`	`object \| None`	`None`	TabPFN regressor used for numerical columns. Defaults to `TabPFNRegressor()` when needed.
`categorical_features`	`list[int] \| None`	`None`	Indices of categorical columns. If `None`, they are inferred from the data.
`use_imputed`	`bool`	`False`	How to treat already-imputed columns when imputing later ones. * `False` (default): every column is imputed using only the original observed data as features, so the result does not depend on column order. * `True`: columns are filled left to right in place, so a column imputed earlier is used as a feature for later columns (chained; order-dependent).

Returns

Type	Description
`np.ndarray`	`np.ndarray`: A new array with all missing values imputed.

Reference

tabpfn

tabpfn-client

tabpfn-extensions

REST API

Imputation and data generation

`TabPFNUnsupervisedModel.generate_synthetic_data`

`TabPFNUnsupervisedModel.impute`

`TabPFNUnsupervisedModel.impute_`

`TabPFNUnsupervisedModel.impute_single_permutation_`

`TabPFNUnsupervisedModel.sample_from_model_prediction_`

`impute_column`

`simple_impute`

​TabPFNUnsupervisedModel.generate_synthetic_data

​TabPFNUnsupervisedModel.impute

​TabPFNUnsupervisedModel.impute_

​TabPFNUnsupervisedModel.impute_single_permutation_

​TabPFNUnsupervisedModel.sample_from_model_prediction_

​impute_column

​simple_impute

`TabPFNUnsupervisedModel.generate_synthetic_data`

`TabPFNUnsupervisedModel.impute`

`TabPFNUnsupervisedModel.impute_`

`TabPFNUnsupervisedModel.impute_single_permutation_`

`TabPFNUnsupervisedModel.sample_from_model_prediction_`

`impute_column`

`simple_impute`