> ## Documentation Index
> Fetch the complete documentation index at: https://docs.priorlabs.ai/llms.txt
> Use this file to discover all available pages before exploring further.

> ## Agent Instructions
> TabPFN-3 and later versions support up to 1,000,000 rows, subject to feature count and checkpoint/API limits.
> For large datasets, use per-estimator subsampling, e.g. inference_config={"SUBSAMPLE_SAMPLES": 100_000}, in TabPFNClassifier or TabPFNRegressor from tabpfn or tabpfn_client.
> Pass raw pandas DataFrames to fit and predict. Categorical strings/categories, raw text, and missing feature values are handled automatically; no manual integer/one-hot encoding, imputation, scaling, or outlier removal is needed.

# Embeddings

> scikit-learn style transformer that extracts TabPFN embeddings.

<Info>
  Looking for usage documentation? Check out [Embeddings](/capabilities/embeddings).
</Info>

## Installation

```bash theme={null}
pip install tabpfn-extensions
```

<div className="python-reference-heading">
  <h2 id="embedding-tabpfn-embedding-tabpfnembedding">
    `TabPFNEmbedding`
  </h2>

  <a className="python-reference-source" href="https://github.com/PriorLabs/tabpfn-extensions/blob/840c15a1848a986b39c85bc17efc61e0e377f983/src/tabpfn_extensions/embedding/tabpfn_embedding.py#L29" aria-label="View source for TabPFNEmbedding"><span aria-hidden="true">\</></span> View source <span aria-hidden="true">↗</span></a>
</div>

scikit-learn style transformer that extracts TabPFN embeddings.

When `n_fold >= 2`, [`fit`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit) produces out-of-fold (OOF) embeddings for the
training data — the robust variant from "A Closer Look at TabPFN v2:
Strength, Limitation, and Extension" ([https://arxiv.org/abs/2502.17361](https://arxiv.org/abs/2502.17361)) —
and then refits a single model on the full training set for use on unseen
data. The OOF embeddings are stored on `train_embeddings_` and returned
by [`fit_transform`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit-transform).

`transform(X)` ALWAYS uses the final, full-data model — it does NOT
return cached OOF embeddings, even when `X` happens to equal the
training set. For OOF embeddings call [`fit_transform`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit-transform) (or read
`train_embeddings_`).

Note on output shape: [`transform`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-transform) returns a 3D array of shape
`(n_estimators, n_samples, embed_dim)`. It is not a drop-in input for
`sklearn.pipeline.Pipeline` / `ColumnTransformer` — those expect 2D
output. Pick an ensemble member (`embeds[0]`) or aggregate across
`axis=0` before passing to a downstream 2D estimator.

```python theme={null}
TabPFNEmbedding(
    n_fold: int = 0,
    *,
    model: TabPFNClassifier | TabPFNRegressor | None = None,
    shuffle: bool = False,
    random_state: int | None = None,
    tabpfn_clf: TabPFNClassifier | None = None,
    tabpfn_reg: TabPFNRegressor | None = None,
)
```

**Parameters**

<div className="python-reference-table">
  | Parameter | Type | Default | Description |
  | - | - | - | - |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding--n-fold" /><code className="python-reference-parameter">n\_<wbr />fold</code> | <code className="python-reference-type">int</code> | `0` | Number of folds for cross-validation. `0` disables CV — the model is trained once on the entire training set and used for both train and unseen data. |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding--model" /><code className="python-reference-parameter">model</code> | <code className="python-reference-type">Tab<wbr />PFN<wbr />Classifier \| Tab<wbr />PFN<wbr />Regressor \| None</code> | `None` | Pre-configured TabPFN estimator. When `None`, the task is inferred from `y` at [`fit`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit) time and a warning is emitted. |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding--shuffle" /><code className="python-reference-parameter">shuffle</code> | <code className="python-reference-type">bool</code> | `False` | Whether to shuffle the K-fold split. Independent of `random_state`. |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding--random-state" /><code className="python-reference-parameter">random\_<wbr />state</code> | <code className="python-reference-type">int \| None</code> | `None` | Seed used by the K-fold split when `shuffle=True`. |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding--tabpfn-clf" /><code className="python-reference-parameter">tabpfn\_<wbr />clf</code> | <code className="python-reference-type">Tab<wbr />PFN<wbr />Classifier \| None</code> | `None` | DEPRECATED. Use `model=` instead. |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding--tabpfn-reg" /><code className="python-reference-parameter">tabpfn\_<wbr />reg</code> | <code className="python-reference-type">Tab<wbr />PFN<wbr />Regressor \| None</code> | `None` | DEPRECATED. Use `model=` instead. |
</div>

**Attributes**

| Attribute | Type | Description |
| - | - | - |
| `model_` | <code className="python-reference-type">Tab<wbr />PFN<wbr />Classifier or Tab<wbr />PFN<wbr />Regressor</code> | The fitted TabPFN model (cloned from `model` or auto-constructed). After [`fit`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit) with `n_fold >= 2` this is the model trained on the full training set. |
| `train_embeddings_` | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Embeddings for the training set. For `n_fold >= 2` these are OOF embeddings aligned to the original sample order; for `n_fold == 0` they are produced by the single full-data model. |

**Examples**

```python theme={null}
>>> from tabpfn_extensions.embedding import TabPFNEmbedding
>>> embedding = TabPFNEmbedding(n_fold=5)
>>> train_embeds = embedding.fit_transform(X_train, y_train)  # OOF
>>> test_embeds = embedding.transform(X_test)                 # final model
```

***

<div className="python-reference-heading">
  <h2 id="embedding-tabpfn-embedding-tabpfnembedding-fit">
    `TabPFNEmbedding.fit`
  </h2>

  <a className="python-reference-source" href="https://github.com/PriorLabs/tabpfn-extensions/blob/840c15a1848a986b39c85bc17efc61e0e377f983/src/tabpfn_extensions/embedding/tabpfn_embedding.py#L177" aria-label="View source for TabPFNEmbedding.fit"><span aria-hidden="true">\</></span> View source <span aria-hidden="true">↗</span></a>
</div>

***

<div className="python-reference-heading">
  <h2 id="embedding-tabpfn-embedding-tabpfnembedding-fit-transform">
    `TabPFNEmbedding.fit_transform`
  </h2>

  <a className="python-reference-source" href="https://github.com/PriorLabs/tabpfn-extensions/blob/840c15a1848a986b39c85bc17efc61e0e377f983/src/tabpfn_extensions/embedding/tabpfn_embedding.py#L218" aria-label="View source for TabPFNEmbedding.fit_transform"><span aria-hidden="true">\</></span> View source <span aria-hidden="true">↗</span></a>
</div>

Fit and return embeddings for the training data.

For `n_fold >= 2` these are out-of-fold embeddings. For
`n_fold == 0` they come from the single full-data model.

```python theme={null}
TabPFNEmbedding.fit_transform(
    X: np.ndarray,
    y: np.ndarray,
) -> np.ndarray
```

**Parameters**

<div className="python-reference-table">
  | Parameter | Type | Default | Description |
  | - | - | - | - |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-fit-transform--x" /><code className="python-reference-parameter">X</code> | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Required | — |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-fit-transform--y" /><code className="python-reference-parameter">y</code> | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Required | — |
</div>

**Returns**

<div className="python-reference-table python-reference-returns">
  | Type | Description |
  | - | - |
  | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | — |
</div>

***

<div className="python-reference-heading">
  <h2 id="embedding-tabpfn-embedding-tabpfnembedding-get-embeddings">
    `TabPFNEmbedding.get_embeddings`
  </h2>

  <a className="python-reference-source" href="https://github.com/PriorLabs/tabpfn-extensions/blob/840c15a1848a986b39c85bc17efc61e0e377f983/src/tabpfn_extensions/embedding/tabpfn_embedding.py#L231" aria-label="View source for TabPFNEmbedding.get_embeddings"><span aria-hidden="true">\</></span> View source <span aria-hidden="true">↗</span></a>
</div>

DEPRECATED. Use [`fit_transform`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit-transform) (OOF) or [`transform`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-transform) (unseen).

```python theme={null}
TabPFNEmbedding.get_embeddings(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X: np.ndarray,
    data_source: str,
) -> np.ndarray
```

**Parameters**

<div className="python-reference-table">
  | Parameter | Type | Default | Description |
  | - | - | - | - |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-get-embeddings--x-train" /><code className="python-reference-parameter">X\_<wbr />train</code> | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Required | — |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-get-embeddings--y-train" /><code className="python-reference-parameter">y\_<wbr />train</code> | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Required | — |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-get-embeddings--x" /><code className="python-reference-parameter">X</code> | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Required | — |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-get-embeddings--data-source" /><code className="python-reference-parameter">data\_<wbr />source</code> | <code className="python-reference-type">str</code> | Required | — |
</div>

**Returns**

<div className="python-reference-table python-reference-returns">
  | Type | Description |
  | - | - |
  | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | — |
</div>

***

<div className="python-reference-heading">
  <h2 id="embedding-tabpfn-embedding-tabpfnembedding-transform">
    `TabPFNEmbedding.transform`
  </h2>

  <a className="python-reference-source" href="https://github.com/PriorLabs/tabpfn-extensions/blob/840c15a1848a986b39c85bc17efc61e0e377f983/src/tabpfn_extensions/embedding/tabpfn_embedding.py#L203" aria-label="View source for TabPFNEmbedding.transform"><span aria-hidden="true">\</></span> View source <span aria-hidden="true">↗</span></a>
</div>

Embed unseen data `X` using the full-data model.

Use this method when you have new, held-out data that was not part of
training.  It always runs inference through `model_` (trained on the
full training set) and never returns cached embeddings.

If you want embeddings for the *training* data, prefer
`fit_transform(X_train, y_train)`, which yields out-of-fold
embeddings for `n_fold >= 2` (avoiding label leakage) or reads
`train_embeddings_` after a [`fit`](/api-reference/python/tabpfn-extensions/embedding#embedding-tabpfn-embedding-tabpfnembedding-fit) call.

```python theme={null}
TabPFNEmbedding.transform(
    X: np.ndarray,
) -> np.ndarray
```

**Parameters**

<div className="python-reference-table">
  | Parameter | Type | Default | Description |
  | - | - | - | - |
  | <span id="embedding-tabpfn-embedding-tabpfnembedding-transform--x" /><code className="python-reference-parameter">X</code> | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | Required | — |
</div>

**Returns**

<div className="python-reference-table python-reference-returns">
  | Type | Description |
  | - | - |
  | <code className="python-reference-type"><a href="https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html">np.ndarray</a></code> | — |
</div>


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.