Metadata-Version: 2.4
Name: tabpfn-v3-mlx
Version: 0.1.0
Summary: TabPFN v3 natively on Apple Silicon via MLX
Project-URL: Homepage, https://github.com/dgallitelli/tabpfn-v3-mlx
Project-URL: Documentation, https://github.com/dgallitelli/tabpfn-v3-mlx#readme
Project-URL: Repository, https://github.com/dgallitelli/tabpfn-v3-mlx
Project-URL: Issues, https://github.com/dgallitelli/tabpfn-v3-mlx/issues
Author-email: Davide Gallitelli <davidegallitelli@gmail.com>
License: MIT
License-File: LICENSE
Keywords: apple-silicon,classification,in-context-learning,mlx,tabpfn,tabular
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: MacOS
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: huggingface-hub>=0.20.0
Requires-Dist: mlx>=0.22.0
Requires-Dist: numpy>=1.24.0
Provides-Extra: convert
Requires-Dist: torch>=2.0.0; extra == 'convert'
Provides-Extra: dev
Requires-Dist: build; extra == 'dev'
Requires-Dist: mypy>=1.0; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: ruff>=0.4.0; extra == 'dev'
Requires-Dist: twine; extra == 'dev'
Provides-Extra: examples
Requires-Dist: scikit-learn>=1.0; extra == 'examples'
Requires-Dist: tabpfn>=8.0; extra == 'examples'
Description-Content-Type: text/markdown

# TabPFN v3 MLX

[![CI](https://github.com/dgallitelli/tabpfn-v3-mlx/actions/workflows/ci.yml/badge.svg)](https://github.com/dgallitelli/tabpfn-v3-mlx/actions/workflows/ci.yml)

Native Apple MLX port of [TabPFN v3](https://github.com/PriorLabs/TabPFN) — the full 53M parameter tabular foundation model running natively on Apple Silicon.

TabPFN v3 performs classification and regression via **in-context learning** — given training data and test features, it produces predictions in a single forward pass with no gradient descent.

This port runs the complete architecture natively on M1/M2/M3/M4/M5 via Apple's [MLX](https://github.com/ml-explore/mlx) framework with zero-copy unified memory.

## Architecture

The full v3 pipeline (2,395 lines in PyTorch) ported to MLX:

```
┌─────────────────────────────────────────────────────────────────────┐
│                     TabPFN v3 Forward Pass                           │
├─────────────────────────────────────────────────────────────────────┤
│                                                                      │
│  Stage 0: Preprocessing                                              │
│  ─────────────────────                                               │
│  x_raw → NaN indicators → mean imputation → z-score scaling         │
│        → circular-shift feature grouping (groups of 3)               │
│                                                                      │
│  Stage 1: Cell + Target Embedding                                    │
│  ────────────────────────────────                                    │
│  grouped_features → Linear(G, 128) → add target embedding (train)   │
│                                                                      │
│  Stage 2a: Distribution Embedding (× 3 blocks)                      │
│  ──────────────────────────────────────────────                      │
│  InducedSelfAttention per column:                                    │
│    inducing_points → cross_attn(ind, train) → hidden                 │
│    all_rows → cross_attn(rows, hidden) → updated embeddings          │
│  Complexity: O(R × n_inducing) instead of O(R²)                      │
│                                                                      │
│  Stage 2b: Column Aggregation (× 3 blocks + readout)                 │
│  ────────────────────────────────────────────────────                 │
│  Prepend CLS tokens → self-attention over features (with RoPE)       │
│  Last block: CLS cross-attends to full sequence → (B, R, 4, 128)    │
│                                                                      │
│  Flatten: (B, R, 4, 128) → (B, R, 512)                              │
│                                                                      │
│  Stage 3: ICL Transformer (× 24 layers)                              │
│  ──────────────────────────────────────                               │
│  Pre-norm RMSNorm → ICL Attention (K/V from train only)              │
│  + SoftmaxScalingMLP (learned query scaling)                         │
│  + GQA (optional fewer KV heads for test rows)                       │
│  + MLP (GELU, no bias, zero-init output)                             │
│  Supports KV caching for efficient repeated inference                │
│                                                                      │
│  Stage 4: Decoder                                                    │
│  ────────────────                                                    │
│  Multiclass: attention retrieval (test→train with one-hot values)    │
│  Regression: MLP → bar distribution buckets                          │
│                                                                      │
│  Output: logits → softmax → probabilities                            │
│                                                                      │
└─────────────────────────────────────────────────────────────────────┘
```

## Installation

```bash
pip install tabpfn-v3-mlx
```

For weight conversion from PyTorch checkpoints:

```bash
pip install "tabpfn-v3-mlx[convert]"
```

## Quick Start

```python
import numpy as np
from tabpfn_mlx import TabPFNV3, TabPFNV3Config, load_v3_pytorch_weights

# Initialize with default config (53M params, 24 ICL layers)
config = TabPFNV3Config(max_num_classes=2)
model = TabPFNV3(config, task_type="multiclass")

# Load pretrained weights (when available)
# model = load_v3_pytorch_weights(model, "path/to/checkpoint.safetensors")

# Predict
probs = model.predict_proba(X_train, y_train, X_test)
preds = model.predict(X_train, y_train, X_test)
```

## Configuration

```python
from tabpfn_mlx import TabPFNV3Config

config = TabPFNV3Config(
    embed_dim=128,                  # Base embedding dimension
    dist_embed_num_blocks=3,        # Distribution embedder layers
    dist_embed_num_heads=8,         # Heads in distribution embedder
    dist_embed_num_inducing_points=128,  # SetTransformer inducing points
    feat_agg_num_blocks=3,          # Column aggregator layers
    feat_agg_num_heads=8,           # Heads in column aggregator
    feat_agg_num_cls_tokens=4,      # CLS tokens (icl_emsize = embed_dim × this)
    nlayers=24,                     # ICL transformer depth
    icl_num_heads=8,                # ICL attention heads
    icl_num_kv_heads=None,          # GQA KV heads (None = standard MHA)
    ff_factor=2,                    # MLP expansion factor
    max_num_classes=10,             # Maximum classes supported
    feature_group_size=3,           # Circular-shift group size
    use_nan_indicators=True,        # NaN/Inf indicator features
)
# icl_emsize = 128 × 4 = 512
```

## KV Cache (Efficient Repeated Inference)

```python
# Build cache from training data (one-time cost)
logits, cache = model(x, y, return_kv_cache=True)

# Reuse cache for new test batches (skips stages 0-2 + K/V projection)
logits_new = model(x_test_only, y, kv_cache=cache, x_is_test_only=True)
```

## Key Differences from nanoTabPFN (v2)

| Aspect | nanoTabPFN | TabPFN v3 |
|--------|-----------|-----------|
| Parameters | 356K | 53M |
| Layers | 3 | 24 ICL + 3 dist + 3 agg |
| Normalization | Post-norm LayerNorm | Pre-norm RMSNorm |
| Feature attention | Direct O(R²) | Induced O(R×k) |
| Positional encoding | None | RoPE + SoftmaxScalingMLP |
| GQA | No | Yes |
| KV cache | No | Multi-level |
| Decoder | MLP | Attention retrieval |

## Development

```bash
git clone https://github.com/dgallitelli/tabpfn-v3-mlx.git
cd tabpfn-v3-mlx
pip install -e ".[dev]"
pytest
```

## Citation

```bibtex
@article{hollmann2025tabpfn,
    title={Accurate Predictions on Small Data with a Tabular Foundation Model},
    author={Hollmann, Noah and Müller, Samuel and Purucker, Lennart and
            Krishnakumar, Arjun and Körfer, Max and Hoo, Shi Bin and
            Schirrmeister, Robin Tibor and Hutter, Frank},
    journal={Nature},
    year={2025}
}
```

## License

MIT. The TabPFN v3 model architecture and weights are subject to [their own license](https://github.com/PriorLabs/TabPFN/blob/main/LICENSE).
