diff --git a/src/tabpfn/architectures/__init__.py b/src/tabpfn/architectures/__init__.py index 4a61e8afb..538ef17de 100644 --- a/src/tabpfn/architectures/__init__.py +++ b/src/tabpfn/architectures/__init__.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING -from . import tabpfn_v2, tabpfn_v2_5, tabpfn_v2_6, tabpfn_v3 +from . import tabpfn_v2, tabpfn_v2_5, tabpfn_v2_6, tabpfn_v3, tabpfn_v3_5 if TYPE_CHECKING: from tabpfn.architectures.interface import ArchitectureModule @@ -23,6 +23,7 @@ "tabpfn_v2_5": tabpfn_v2_5, "tabpfn_v2_6": tabpfn_v2_6, "tabpfn_v3": tabpfn_v3, + "tabpfn_v3_5": tabpfn_v3_5, } """Map from architecture names to the corresponding module.""" diff --git a/src/tabpfn/architectures/tabpfn_v3_5.py b/src/tabpfn/architectures/tabpfn_v3_5.py new file mode 100644 index 000000000..ecc6f4217 --- /dev/null +++ b/src/tabpfn/architectures/tabpfn_v3_5.py @@ -0,0 +1,3259 @@ +# ruff: noqa: PLR0912, C901 +"""TabPFN v3.5 architecture, inference only. + +`task_type` is a per-`forward()` argument, so one model instance handles both +multiclass and regression. + +Shape suffix convention: + +B: batch size +R: total rows (train + test) +Ri: input rows, could be either train + test or test. +Rj: Chunked rows (<= R) +N: train rows +M: test rows +C: total columns +Cj: Chunked columns (<= C) +E: embedding dimension +T: Target dim (e.g. number of classes). +Cl: number of CLS tokens + +D: head dimension +H: num heads +S: sequence length + +Copyright (c) Prior Labs GmbH 2026. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import logging as _logging +import math +from collections.abc import Callable +from functools import partial +from typing import TYPE_CHECKING, Any, Literal, cast +from typing_extensions import override + +import numpy as np +import pydantic +import torch +import torch.nn.functional as F # noqa: N812 +import torch.utils.checkpoint +from torch import nn + +from tabpfn.architectures.interface import ( + Architecture, + ArchitectureConfig, + PerformanceOptions, +) +from tabpfn.architectures.kv_cache import ( + QUANTIZED_KV_DTYPE, + KVCache, + KVCacheEntry, + QuantizedKVCacheEntry, +) +from tabpfn.architectures.shared.chunked_evaluate import chunked_evaluate_maybe_inplace +from tabpfn.architectures.shared.scaled_dot_product_attention import ( + scaled_dot_product_attention, +) +from tabpfn.errors import is_oom_error +from tabpfn.preprocessing.torch.torch_standard_scaler import TorchStandardScaler + +if TYPE_CHECKING: + from torch.nn.attention import SDPBackend + + from tabpfn.constants import TaskType + + +_logger = _logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass +class TabPFNV3p5Config(ArchitectureConfig): + """Configuration for the single-file TabPFN v3.5 architecture. + + The defaults are the v3.5 pre-release checkpoint's config, so the only keys a + caller has to supply are the head sizes `max_num_classes` and `num_buckets`, + which the checkpoint carries. + """ + + name: str = "TabPFN-v3.5" + + # ---- Distribution embedder (per-column induced self-attention) ---- + embed_dim: int = 128 + """Base embedding dimension used throughout the model.""" + + dist_embed_num_blocks: int = 3 + """Number of induced-self-attention blocks in the distribution embedder.""" + + dist_embed_num_heads: int = 8 + """Number of attention heads in the distribution embedder.""" + + dist_embed_num_inducing_points: int = 128 + """Number of inducing points in the distribution embedder.""" + + feature_group_size: int = 3 + """Number of features per circular-shift group in the distribution embedder.""" + + # ---- Feature aggregation (cross-feature interaction via CLS tokens) ---- + feat_agg_num_blocks: int = 3 + """Number of transformer blocks in the feature aggregation stage.""" + + feat_agg_num_heads: int = 8 + """Number of attention heads in the feature aggregation stage.""" + + feat_agg_num_cls_tokens: int = 8 + """Number of CLS tokens used to aggregate per-row feature information.""" + + feat_agg_rope_base: float = 100_000 + """RoPE base in the feature aggregation transformer.""" + + # ---- ICL transformer ---- + nlayers: int = 24 + """Number of transformer blocks in the ICL stage.""" + + icl_num_heads: int = 16 + """Number of attention heads in the ICL stage.""" + + icl_num_kv_heads: int | None = None + """GQA: number of KV heads in the ICL stage. None = standard MHA. + Must divide icl_num_heads.""" + + icl_num_kv_heads_test: int | None = 1 + """Number of KV heads used by test rows in the ICL stage. + None = same as train rows (i.e. icl_num_kv_heads / standard MHA). + Any value that divides icl_num_heads is valid (1 = MQA, other = GQA).""" + + # ---- Output decoder (many-class for multiclass, MLP for regression) ---- + decoder_head_dim: int = 64 + """Head dimension for the many-class decoder attention.""" + + decoder_num_heads: int = 6 + """Number of attention heads for the many-class decoder.""" + + decoder_use_softmax_scaling: bool = True + """If True, apply softmax scaling in the many-class decoder.""" + + # ---- Shared ---- + ff_factor: int = 2 + """Feed-forward expansion factor used throughout the model.""" + + softmax_scaling_mlp_hidden_dim: int = 64 + """Number of hidden units in the MLPs for the SoftmaxScalingMLP layer.""" + + # ---- Fourier cell embedding ---- + fourier_encoding_num_frequencies: int = 32 + """Number of learnable Fourier frequencies per grouped cell value. Each value + channel is expanded into twice this many sin/cos features.""" + + cell_ecdf_num_frequencies: int = 4 + """Number of Fourier frequencies for the per-cell ECDF channel; each cell + contributes twice this many metadata features.""" + + cell_ecdf_num_buckets: int = 8192 + """Number of bucket edges the per-cell ECDF ranks against, per column. + + Caps the ECDF context at `num_buckets` values per column instead of one per + train row, which is what keeps the inference cache from growing with the + table. A column with at most this many distinct values is ranked exactly; + above it, ranks between two edges are interpolated. Set it to at least the + train-row count to rank every table exactly.""" + + cell_embed_row_chunk_size: int | None = 2048 + """Row-chunk size for the Fourier cell embedder. When set, the embedder splits + the row axis into chunks of this size, bounding the peak memory of the Fourier + expansion's `(..., G, 2F)` features to one chunk. Matters most on tall tables, + where the per-column inducing-hidden pass embeds all train rows at once. None + disables chunking. Ignored under torch.compile, which plans its own + recomputation.""" + + # ---- Memory-efficient inference ---- + inference_row_chunk_size: int = 2048 + """Max rows per Stage 0-2 chunk during inference.""" + + inference_col_chunk_size: int = 4 + """Max output groups per chunk for inducing hidden state computation.""" + + def __post_init__(self) -> None: + """Validate config constraints.""" + for name in ( + "fourier_encoding_num_frequencies", + "cell_ecdf_num_frequencies", + ): + if getattr(self, name) < 1: + raise ValueError(f"{name} must be >= 1, got {getattr(self, name)}") + # A single bucket leaves no interval to interpolate over. + if self.cell_ecdf_num_buckets < 2: + raise ValueError( + f"cell_ecdf_num_buckets must be >= 2, got {self.cell_ecdf_num_buckets}" + ) + if ( + self.cell_embed_row_chunk_size is not None + and self.cell_embed_row_chunk_size <= 0 + ): + raise ValueError( + "cell_embed_row_chunk_size must be > 0 or None, got " + f"{self.cell_embed_row_chunk_size}" + ) + if self.icl_num_kv_heads is not None and ( + self.icl_num_heads % self.icl_num_kv_heads != 0 + ): + raise ValueError( + f"icl_num_heads ({self.icl_num_heads}) must be divisible by " + f"icl_num_kv_heads ({self.icl_num_kv_heads})" + ) + if self.icl_num_kv_heads_test is not None: + if self.icl_num_heads % self.icl_num_kv_heads_test != 0: + raise ValueError( + f"icl_num_heads ({self.icl_num_heads}) must be divisible by " + f"icl_num_kv_heads_test ({self.icl_num_kv_heads_test})" + ) + effective_kv = ( + self.icl_num_kv_heads + if self.icl_num_kv_heads is not None + else self.icl_num_heads + ) + if self.icl_num_kv_heads_test > effective_kv: + raise ValueError( + f"icl_num_kv_heads_test ({self.icl_num_kv_heads_test}) must be " + f"<= the number of train KV heads ({effective_kv})" + ) + + +# --------------------------------------------------------------------------- +# TabPFN v3.5 KV cache +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class TabPFNV3p5Cache(KVCache): + """Top-level cache container for the TabPFN v3.5 explicit KV cache. + + Stores everything needed to skip stages 0-2 for train rows and reuse + cached K/V in the ICL transformer. + + Attributes: + kv: Per-layer KV cache for the ICL transformer blocks. + decoder_keys: Projected many-class decoder keys of shape + `(B, N_train, H_dec, D_dec)`, i.e. the decoder's `k_projection` + already applied to the post-ICL, post-norm train embeddings. `None` + for a regression cache, which has no many-class decoder. Caching the + keys rather than the embeddings they come from is smaller + (`H_dec * D_dec` is below the ICL width) and keeps the projection off + the predict path. + train_shape: `(batch_size, num_train)` for validation. + scaler_cache: Fitted standard-scaler statistics (`mean`, `std`). Allows + standardising test-only data without train rows present. + ecdf_context: ECDF bucket edges and their rank bounds per column, `(3, + B, C, K)` at `ECDF_CONTEXT_DTYPE`, against which test cells are + ranked. `K` is `min(cell_ecdf_num_buckets, n_train)`, so on a tall + table this stops growing with the row count. + inducing_hidden: Per-block inducing hidden states from the + distribution embedder, each of shape `(B*C_out, n_ind, E)`. + Allows running `cross_attn_block2` on test rows without + recomputing `cross_attn_block1` from train rows. + """ + + decoder_keys: torch.Tensor | None = None + train_shape: tuple[int, int] = (0, 0) + scaler_cache: dict[str, torch.Tensor] | None = None + ecdf_context: torch.Tensor | None = None + inducing_hidden: list[torch.Tensor] | None = None + + @override + def to(self, device: torch.device | str) -> TabPFNV3p5Cache: + """Move all cached tensors to the given device.""" + return TabPFNV3p5Cache( + kv=self._kv_to(device), + decoder_keys=( + self.decoder_keys.to(device) if self.decoder_keys is not None else None + ), + train_shape=self.train_shape, + scaler_cache=self._dict_of_tensors_to(self.scaler_cache, device), + ecdf_context=( + self.ecdf_context.to(device) if self.ecdf_context is not None else None + ), + inducing_hidden=self._list_of_tensors_to(self.inducing_hidden, device), + ) + + def quantize(self, dtype: torch.dtype = QUANTIZED_KV_DTYPE) -> TabPFNV3p5Cache: + """Return a new cache with quantized ICL KV entries. + + Only the ICL KV cache is quantized; `decoder_keys`, `scaler_cache`, + `ecdf_context` and `inducing_hidden` stay at the precision they were built + at. `InferenceEngineExplicitKVCache` + calls this whenever the resolved `kv_cache_precision` is not `"auto"`, which + is why `get_supported_kv_cache_precisions` has to advertise the dtypes this + handles. + + Args: + dtype: Target quantization dtype (default `QUANTIZED_KV_DTYPE`, int8; + `FP8_KV_DTYPE` is the other one the engine can ask for). + """ + quantized_kv = { + idx: (entry.quantize(dtype) if isinstance(entry, KVCacheEntry) else entry) + for idx, entry in self.kv.items() + } + return TabPFNV3p5Cache( + kv=quantized_kv, + decoder_keys=self.decoder_keys, + train_shape=self.train_shape, + scaler_cache=self.scaler_cache, + ecdf_context=self.ecdf_context, + inducing_hidden=self.inducing_hidden, + ) + + +def get_cache_size( + *, + n_train: int, + n_features: int, + model_config: TabPFNV3p5Config, + task_type: TaskType, + base_dtype: torch.dtype | Literal["autocast"], + kv_cache_precision: Literal["auto", "int8", "fp8"] = "int8", +) -> int: + """Cached memory in bytes for a single TabPFN v3.5 estimator at batch size 1. + + Works from shapes alone, so it can be called before fitting to size an + inference run. It is the exact resident size of one estimator's + `TabPFNV3p5Cache`, summing every tensor the cache holds: + + 1. The ICL transformer KV cache (int8 plus per-tensor scales when quantized). + 2. The many-class decoder keys, `(n_train, H_dec * D_dec)`. Multiclass only — + regression has no many-class decoder and caches nothing here, which is why + this needs the `task_type` the forward pass will be called with. + 3. The distribution-embedder `inducing_hidden` states. + 4. The fitted scaler stats, `mean` and `std`. + 5. The ECDF ranking context `ecdf_context`, always at `ECDF_CONTEXT_DTYPE`, + whatever the compute precision. + + The cache is not uniformly one dtype, so each term is sized at its own + precision, selected by `base_dtype`: + + * **Forced precision** (a `torch.dtype`, mirroring `inference_precision` set + to a dtype): the model and inputs are cast to it, so every non-KV term + lands at that dtype. + * **Autocast** (`"autocast"`, the GPU default for `inference_precision="auto"`): + weights stay fp32 and ops are cast at runtime to fp16. The matmul-lineage + tensors (KV, and `decoder_keys` via its explicit cast to the KV dtype) + take fp16, while the reduction/norm-lineage tensors (`inducing_hidden`, + `scaler_cache`) stay fp32. + + Args: + n_train: Number of training rows. The KV cache and the decoder keys + scale with this; test rows are not cached. + n_features: Number of feature columns the model sees. Exact for the + columns the model sees; for real end-to-end runs preprocessing may + change it (SVD features, categorical expansion, per-member + subsampling), making those terms approximate. + model_config: The v3.5 architecture config. + task_type: The task the cache will be built for. Selects whether the + many-class decoder keys are counted. + base_dtype: A `torch.dtype` for the forced-precision path, or + `"autocast"` for the GPU autocast path. + kv_cache_precision: If `"int8"` (default) or `"fp8"`, the KV cache is + sized at one byte per element plus per-tensor scales at the KV + compute dtype, mirroring the engine's `maybe_quantize_kv_cache`; if + `"auto"`, the K/V are sized at the compute dtype with no scales. + + Returns: + Per-estimator cache size in bytes. Multiply by the ensemble size for the + total (each estimator holds its own cache). + """ + if kv_cache_precision not in ("auto", "int8", "fp8"): + raise ValueError( + f"Invalid kv_cache_precision: {kv_cache_precision}. " + "Must be one of 'auto', 'int8' or 'fp8'." + ) + quantize_kv_cache = kv_cache_precision in ("int8", "fp8") + + if base_dtype == "autocast": + kv_dtype = QUANTIZED_KV_DTYPE if quantize_kv_cache else torch.float16 + kv_scale_dtype = torch.float16 # per-tensor scales, at the KV fp16 dtype + decoder_key_dtype = torch.float16 + inducing_dtype = torch.float32 + scaler_dtype = torch.float32 + else: + kv_dtype = QUANTIZED_KV_DTYPE if quantize_kv_cache else base_dtype + kv_scale_dtype = base_dtype + decoder_key_dtype = base_dtype + inducing_dtype = base_dtype + scaler_dtype = base_dtype + + icl_emsize = model_config.embed_dim * model_config.feat_agg_num_cls_tokens + head_dim = icl_emsize // model_config.icl_num_heads + if model_config.icl_num_kv_heads_test is not None: + num_kv_heads = model_config.icl_num_kv_heads_test + elif model_config.icl_num_kv_heads is not None: + num_kv_heads = model_config.icl_num_kv_heads + else: + num_kv_heads = model_config.icl_num_heads + + # 1. ICL KV cache: key + value (the factor of 2), per layer, over all layers. + kv_elements = model_config.nlayers * 2 * n_train * num_kv_heads * head_dim + total_bytes = kv_elements * kv_dtype.itemsize + if quantize_kv_cache: + # One scalar scale per key and per value tensor. + total_bytes += model_config.nlayers * 2 * kv_scale_dtype.itemsize + + # 2. Many-class decoder keys, (n_train, H_dec * D_dec). Multiclass only. + if task_type == "multiclass": + decoder_key_width = ( + model_config.decoder_num_heads * model_config.decoder_head_dim + ) + total_bytes += n_train * decoder_key_width * decoder_key_dtype.itemsize + + # 3. Distribution-embedder inducing states: one + # (n_features, dist_embed_num_inducing_points, embed_dim) tensor per block. + total_bytes += ( + model_config.dist_embed_num_blocks + * n_features + * model_config.dist_embed_num_inducing_points + * model_config.embed_dim + ) * inducing_dtype.itemsize + + # 4. Fitted scaler stats: mean + std, each (n_features,). + total_bytes += 2 * n_features * scaler_dtype.itemsize + + # 5. The ECDF ranking context: an edge value and its two rank bounds (the + # factor of 3) per bucket, per column. + num_buckets = min(model_config.cell_ecdf_num_buckets, n_train) + total_bytes += 3 * n_features * num_buckets * ECDF_CONTEXT_DTYPE.itemsize + + return total_bytes + + +# --------------------------------------------------------------------------- +# Rotary Positional Embeddings (RoPE) — compile-friendly, no einops +# --------------------------------------------------------------------------- +# We don't cache cos/sin, since this blocks torch.compile. + + +def apply_rope( + t: torch.Tensor, + inv_freq: torch.Tensor, + *, + interleaved: bool = False, +) -> torch.Tensor: + """Apply rotary positional embeddings to `t` along seq_dim=-2. + + All intermediate math is done in `inv_freq.dtype` (fp32 by + construction) and the result is cast back to `t.dtype`. + + Args: + t: Tensor of shape `(..., S, D)` where the head dim `D` is + even. The sequence dim is the second-to-last axis. + inv_freq: `(D // 2,)` inverse frequencies (typically + `1 / theta ** (2i / D)`). + interleaved: When `True`, rotates dimension pairs + `(0, 1), (2, 3), …` (LLaMA/HF interleaved layout). When + `False` (default), splits the last dim into two contiguous + halves and rotates them against each other. + """ + dtype = t.dtype + seq_len = t.shape[-2] + positions = torch.arange(seq_len, device=t.device, dtype=inv_freq.dtype) + freqs = positions[:, None] * inv_freq[None, :] # (S, D/2) + cos = freqs.cos() + sin = freqs.sin() + if interleaved: + cos = cos.repeat_interleave(2, dim=-1) # (S, D) + sin = sin.repeat_interleave(2, dim=-1) + t_even = t[..., 0::2] + t_odd = t[..., 1::2] + # stack → (..., D/2, 2) rows (-t_odd, t_even); flatten → (-t1, t0, -t3, t2, …) + t_rotated = torch.stack((-t_odd, t_even), dim=-1).flatten(-2) + else: + cos = torch.cat((cos, cos), dim=-1) # (S, D) + sin = torch.cat((sin, sin), dim=-1) + half = t.shape[-1] // 2 + t_rotated = torch.cat((-t[..., half:], t[..., :half]), dim=-1) + return (t * cos + t_rotated * sin).to(dtype) + + +class RotaryEmbedding(nn.Module): + """Compile-friendly rotary positional embedding. + + Args: + dim: Per-head rotation dimension. Must be even. + theta: Base for the rotary frequencies (10_000 in the original + paper, 100_000 in our configs). + interleaved: See `apply_rope`. + """ + + def __init__( + self, + dim: int, + *, + theta: float = 10_000.0, + interleaved: bool = False, + ) -> None: + super().__init__() + assert dim % 2 == 0, f"RoPE head dim must be even, got {dim}" + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + # Store as a non-learnable nn.Parameter (not a buffer) to match the + # upstream RotaryEmbedding which has `self.freqs = nn.Parameter(..., + # requires_grad=False)`. This preserves the parameter count seen by + # the optimizer, avoiding subtle numerical drift in training due to + # Adam state ordering changes. + self.freqs = nn.Parameter(inv_freq, requires_grad=False) + self.interleaved = interleaved + + def rotate_queries_or_keys(self, t_BHSD: torch.Tensor) -> torch.Tensor: + """Apply RoPE to t_BSHD.""" + return apply_rope(t_BHSD, self.freqs, interleaved=self.interleaved) + + +class _DtypeMatchingRMSNorm(nn.RMSNorm): + """RMSNorm that casts weight to match the input dtype. + + Fused CUDA kernels require matching dtypes; casting the tiny weight/bias per-call + avoids unfused fallbacks under autocast. + """ + + @override + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.weight.dtype != input.dtype: + return F.rms_norm( + input, + self.normalized_shape, + self.weight.to(input.dtype), + self.eps, + ) + return super().forward(input) + + +class ManyClassDecoder(nn.Module): + """Attention-based retrieval decoder for many-class classification. + + Computes weighted (by attention score) average over one-hot encoded + train targets, then takes the log to obtain logits. Supports arbitrary + class counts by chunking the value (one-hot) dimension into head_dim-sized + pieces and folding them into the batch dimension for a single flash-attention + call. + """ + + def __init__( + self, + max_num_classes: int, + input_size: int, + head_dim: int = 64, + num_heads: int = 6, + softmax_scaling_layer: nn.Module | None = None, + ): + """Init.""" + super().__init__() + self.max_num_classes = max_num_classes + self.input_size = input_size + self.attention_size = head_dim * num_heads + self.head_dim = head_dim + self.num_heads = num_heads + self.q_projection = nn.Linear(self.input_size, self.attention_size) + self.k_projection = nn.Linear(self.input_size, self.attention_size) + self.softmax_scaling_layer = softmax_scaling_layer + + def project_keys(self, train_embeddings_BNE: torch.Tensor) -> torch.Tensor: + """Project train embeddings to per-head keys: `(B,N,E)` -> `(B,N,H,D)`. + + Split out from `forward` so the inference cache can hold the keys instead + of the embeddings they come from. That is smaller — `H*D` is below + `input_size` — and keeps the projection off the predict path. + """ + k_BNE = self.k_projection(train_embeddings_BNE) + return k_BNE.view(*k_BNE.shape[:2], self.num_heads, self.head_dim).contiguous() + + @override + def forward( + self, + train_keys_BNHD: torch.Tensor, + test_embeddings_BME: torch.Tensor, + targets_BN: torch.Tensor, + *, + num_present_classes: int, + ) -> torch.Tensor: + """Perform a forward pass, on keys already built by `project_keys`.""" + B, M, _ = test_embeddings_BME.shape + q_BME = self.q_projection(test_embeddings_BME) + # Mirrors the dtype guard in ICLAttention's cached path: keys built under + # autocast, or read back from the cache, may not match the query dtype. + if train_keys_BNHD.dtype != q_BME.dtype: + train_keys_BNHD = train_keys_BNHD.to(q_BME.dtype) + + if M == 0: + # Flash attention rejects a query sequence of length 0, so return + # early. The zero-weighted sums keep both inputs in the graph. + empty = test_embeddings_BME.new_empty((0, B, self.max_num_classes)) + return empty + (q_BME.sum() + train_keys_BNHD.sum()) * 0.0 + + # Mask out non-finite target rows. Those shouldn't contribute to the output + # and the .long conversion results in different values on different platforms. + is_finite_BN = torch.isfinite(targets_BN) + targets_long = torch.where(is_finite_BN, targets_BN.long(), 0) + # Only the classes present in the batch need a one-hot column; the rest + # are zero everywhere and are restored by the padding below. Narrowing + # the class axis shrinks the int64 one-hot and, since + # `_chunked_class_attention` runs `ceil(T / head_dim)` folded attention + # passes, can cut the attention cost by that factor. + one_hot_targets_BNT = torch.where( + is_finite_BN[..., None], + F.one_hot(targets_long, num_classes=num_present_classes), + 0, + ).to(dtype=q_BME.dtype) + + q_BMHD = q_BME.view(B, M, self.num_heads, self.head_dim).contiguous() + k_BNHD = train_keys_BNHD + one_hot_targets_BNHT = ( + one_hot_targets_BNT.unsqueeze(2) + .expand(-1, -1, self.num_heads, -1) + .contiguous() + ) + test_output_BMHT = _chunked_class_attention( + q_BMHD, + k_BNHD, + one_hot_targets_BNHT, + softmax_scaling_layer=self.softmax_scaling_layer, + ) + test_output_BMT = test_output_BMHT.mean(2) # average over heads + + # Widen back to the architectural class count. A class absent from the + # train targets holds an all-zero value column, so attention returns + # exactly the zero written here. + missing_classes = self.max_num_classes - num_present_classes + if missing_classes: + test_output_BMT = F.pad(test_output_BMT, (0, missing_classes)) + + test_output_MBT = test_output_BMT.transpose(0, 1) + # convert to logits: + return torch.log(torch.clamp(test_output_MBT, min=1e-5) + 3e-5) + + +def _chunked_class_attention( + q_BSHD: torch.Tensor, + k_BJHD: torch.Tensor, + v_BJHT: torch.Tensor, + softmax_scaling_layer: nn.Module | None = None, +) -> torch.Tensor: + """Run retrieval attention where the value dimension C may exceed head_dim D. + + Splits V into head_dim-sized chunks along the class axis, folds the chunk + index into the batch dimension, and dispatches a single flash-attention call. + This avoids the O(N*M) memory cost of the math backend for any class count. + + Args: + q_BSHD: Query tensor of shape (B, S, H, D) for test points. + k_BJHD: Key tensor of shape (B, J, H, D) for train points. + v_BJHT: Value tensor of shape (B, J, H, T) holding one-hot class + encodings; T may be larger than D. + softmax_scaling_layer: Optional scaling module to scale queries before SDPA. + + Returns: + Output tensor of shape (B, S, H, T). + """ + B, S, H, D = q_BSHD.shape + T = v_BJHT.shape[-1] + num_chunks = math.ceil(T / D) + + # Pad V to a multiple of D along the class axis + pad = num_chunks * D - T + if pad > 0: + v_BJHT = F.pad(v_BJHT, (0, pad)) + + # Fold chunk index into batch dimension + J = v_BJHT.shape[1] + v_folded = ( + v_BJHT.reshape(B, J, H, num_chunks, D) + .permute(0, 3, 1, 2, 4) + .reshape(B * num_chunks, J, H, D) + .contiguous() + ) + q_folded = ( + q_BSHD.unsqueeze(1) + .expand(-1, num_chunks, -1, -1, -1) + .reshape(B * num_chunks, S, H, D) + .contiguous() + ) + k_folded = ( + k_BJHD.unsqueeze(1) + .expand(-1, num_chunks, -1, -1, -1) + .reshape(B * num_chunks, J, H, D) + .contiguous() + ) + + # Single flash-attention call across all chunks + out_folded = _batched_scaled_dot_product_attention( + q_folded, k_folded, v_folded, softmax_scaling_layer=softmax_scaling_layer + ) + + # Unfold and trim padding: (B*K, S, H, D) -> (B, S, H, T) + return ( + out_folded.reshape(B, num_chunks, S, H, D) + .permute(0, 2, 3, 1, 4) + .reshape(B, S, H, num_chunks * D)[..., :T] + ) + + +class TrainableOrthogonalEmbedding(nn.Module): + """Trainable class embeddings initialized with orthogonal initialization.""" + + def __init__(self, num_classes: int, embed_dim: int) -> None: + super().__init__() + self.embedding = nn.Embedding(num_classes, embed_dim) + self._init() + + @override + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Map integer labels (B, T) -> embeddings (B, T, embed_dim).""" + return self.embedding(x.long()) + + def _init(self) -> None: + """Initialize embedding weight rows orthogonally in-place. + + The first `min(num_classes, embed_dim)` rows are set to orthonormal + vectors via QR decomposition; remaining rows (when `num_classes > + embed_dim`) are unit-normalized random vectors. + """ + weight = self.embedding.weight + num_classes, embed_dim = weight.shape + k = min(num_classes, embed_dim) + q, _ = torch.linalg.qr(torch.randn(embed_dim, k)) + ortho_rows = q.T # (k, embed_dim) + with torch.no_grad(): + weight[:k].copy_(ortho_rows) + if num_classes > embed_dim: + extra = torch.randn(num_classes - k, embed_dim) + extra = extra / extra.norm(dim=-1, keepdim=True).clamp(min=1e-8) + weight[k:].copy_(extra) + + +class MLP(nn.Sequential): + """Two-layer GELU feed-forward network with zero-initialized output.""" + + def __init__( + self, + emsize: int, + dim_feedforward: int, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + kw = {"device": device, "dtype": dtype} + linear2 = nn.Linear(dim_feedforward, emsize, bias=False, **kw) + nn.init.zeros_(linear2.weight) + super().__init__( + nn.Linear(emsize, dim_feedforward, bias=False, **kw), + nn.GELU(), + linear2, + ) + + +class FourierFeatureGroupEmbedder(nn.Module): + """Fourier-feature embedding of a grouped feature block (TabFM-style). + + Maps `(..., G) -> (..., E)`: expands each grouped scalar into `[sin, cos]` + against one learnable frequency bank shared by every cell, sums the features + over the group, and projects with a shared linear. Summing before the linear + (`Σ_g W·f_g = W·Σ_g f_g`) avoids materializing the per-group `(..., G, E)` + projection. The frequency multiply and sin/cos run in fp32; the projection + follows autocast. + """ + + def __init__( + self, + group_size: int, + embed_dim: int, + num_freq: int, + *, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + kw = {"device": device, "dtype": dtype} + # Bias-free: the summed embedding is LayerNorm'd (with a learnable bias) + # in FourierPlusMetadataFeatureGroupEmbedder, so a per-linear bias is + # redundant. + self.frequencies = nn.Parameter(torch.randn(group_size, num_freq, **kw) * 2.0) + self.in_linear = nn.Linear(num_freq * 2, embed_dim, bias=False, **kw) + + @override + def forward(self, x_G: torch.Tensor) -> torch.Tensor: + """Embed grouped cell values `(..., G)` into `(..., E)`.""" + dt = x_G.dtype + proj = x_G.unsqueeze(-1).float() * self.frequencies.float() # (..., G, F) + feats_G = torch.cat([proj.sin(), proj.cos()], dim=-1).to(dt) # (..., G, 2F) + return self.in_linear(feats_G.sum(dim=-2)) # (..., E) + + +class FourierPlusMetadataFeatureGroupEmbedder(nn.Module): + """Fourier-embed the grouped cell values; linearly embed the metadata; sum. + + Maps a grouped cell tensor `(..., G) -> (..., E)`, a drop-in replacement for + the linear cell embedder. The leading `group_size` channels are the + standard-scaled values, which both paths read; the input continues with the + NaN indicators (when enabled) and the trailing `group_size` raw ECDF ranks. + + Two E-dim embeddings are summed: + - the values through a `FourierFeatureGroupEmbedder` (frequencies in fp32, + projection in bf16 under autocast); + - the whole metadata block through one bias-free linear run in fp32, because + bf16's 8-bit mantissa cannot keep high-cardinality ordinal categorical + values distinct. + + A final `LayerNorm` normalizes the summed embedding to unit scale — matching + the LayerNorm applied to the target-aware y-encoder output it is later summed + with, and keeping the scale stable across `group_size` and `num_freq`. + """ + + def __init__( + self, + group_size: int, + embed_dim: int, + num_freq: int, + *, + ecdf_num_frequencies: int, + row_chunk_size: int | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + kw = {"device": device, "dtype": dtype} + self.group_size = group_size + self.ecdf_num_frequencies = ecdf_num_frequencies + self.row_chunk_size = row_chunk_size + self.fourier = FourierFeatureGroupEmbedder( + group_size, embed_dim, num_freq, **kw + ) + # Each grouped value arrives with its NaN/Inf indicator. + value_width = group_size * 2 + # One raw ECDF rank per group position on the way in; each is lifted to + # 2K sin/cos features inside `_embed`, so the linear is wider than the + # tensor the caller passes. + self.input_width = value_width + group_size + self.metadata_width = value_width + group_size * 2 * ecdf_num_frequencies + self.metadata_linear = nn.Linear( + self.metadata_width, embed_dim, bias=False, **kw + ) + self.layernorm = nn.LayerNorm(embed_dim, elementwise_affine=True, **kw) + + @override + def forward(self, x_grouped_G: torch.Tensor) -> torch.Tensor: + """Embed a grouped cell tensor `(B, R, C, G)` into `(B, R, C, E)`.""" + # The Fourier expansions transiently materialize (..., G, 2F) and + # (..., G, 2K) features per cell across the whole row axis (dim=1) at + # once. Chunking the rows bounds that peak to `row_chunk_size` rows; + # rows are independent here, so chunk-and-concat is exact. This matters + # at scale: the per-column inducing-hidden pass embeds all train rows at + # once (only columns are chunked there), so a tall table would otherwise + # blow up the transient along the row axis. + # Skipped under torch.compile: inductor plans its own recomputation, and + # a Python row loop over a dynamic row count would force graph breaks. + if ( + self.row_chunk_size is None + or torch.compiler.is_compiling() + or x_grouped_G.shape[1] <= self.row_chunk_size + ): + return self._embed(x_grouped_G) + parts = [ + self._embed(x_grouped_G[:, start : start + self.row_chunk_size]) + for start in range(0, x_grouped_G.shape[1], self.row_chunk_size) + ] + return torch.cat(parts, dim=1) + + def _embed(self, x_grouped_G: torch.Tensor) -> torch.Tensor: + dt = x_grouped_G.dtype + fourier_out = self.fourier(x_grouped_G[..., : self.group_size]) # (..., E) + # Slice the trailing raw ECDF ranks out of the metadata block and put them + # back as 2K sin/cos features each, which is why the linear is wider than + # the tensor that arrives. Lifting here rather than before grouping bounds + # the (..., G, 2K) expansion by the row chunk and keeps grouping to one + # channel per group position instead of 2K. + ranks_G = x_grouped_G[..., -self.group_size :] + metadata_G = torch.cat( + [ + x_grouped_G[..., : -self.group_size], + _ecdf_fourier_features(ranks_G, self.ecdf_num_frequencies).flatten(-2), + ], + dim=-1, + ) + # fp32 metadata projection even under bf16 autocast. Cast the weight too + # so this holds under bf16 parameters, not just bf16 inputs. + with torch.autocast(device_type=x_grouped_G.device.type, enabled=False): + metadata_out = F.linear( + metadata_G.float(), self.metadata_linear.weight.float() + ) # (..., E) + return self.layernorm((fourier_out.float() + metadata_out).to(dt)) + + +class SoftmaxScalingMLP(nn.Module): + """Query-aware attention scaling using MLPs to compute scaling factors. + + Applies scaling to queries: + + q_scaled = q * base_mlp(logn) * (1 + tanh(query_mlp(q))), + + where the base MLP learns length-dependent scaling and the query MLP + learns query-dependent modulation. + """ + + def __init__( + self, + num_heads: int, + head_dim: int, + n_hidden: int = 64, + ): + """Initializes the SoftmaxScalingMLP module. + + Args: + num_heads: Number of attention heads. + head_dim: Dimension of each attention head. + n_hidden: Number of hidden units in the MLPs. + """ + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + base_out_dim = num_heads * head_dim + query_out_dim = head_dim + + self.base_mlp = nn.Sequential( + nn.Linear(1, n_hidden), nn.GELU(), nn.Linear(n_hidden, base_out_dim) + ) + self.query_mlp = nn.Sequential( + nn.Linear(head_dim, n_hidden), nn.GELU(), nn.Linear(n_hidden, query_out_dim) + ) + # ensures initial modulation is zero + nn.init.zeros_(self.query_mlp[-1].weight) # type: ignore + nn.init.zeros_(self.query_mlp[-1].bias) # type: ignore + + @override + def forward(self, q_BSHD: torch.Tensor, n: int) -> torch.Tensor: + """Applies scalable attention scaling to queries. + + Args: + q_BSHD: Query tensor after projection, shape `[B, S, H, D]`. + B: Batch size. + S: Sequence length. + H: Number of heads. + D: Head dimension. + n: Number of elements for log-n scaling. + + Returns: + Scaled query tensor, same shape as `q_BSHD`. + """ + logn_11 = _safe_log_seqlen(n, q_BSHD.device, q_BSHD.dtype).reshape(1, 1) + base_scales = self.base_mlp(logn_11).view(1, 1, self.num_heads, self.head_dim) + modulation = 1 + torch.tanh(self.query_mlp(q_BSHD)) + scales = base_scales * modulation + return q_BSHD * scales + + +def _batched_scaled_dot_product_attention( + q_BSHD: torch.Tensor, + k_BSJD: torch.Tensor | None, + v_BSJD: torch.Tensor | None, + softmax_scaling_layer: nn.Module | None = None, + _backends_override: list[SDPBackend] | None = None, + quantized_kv: QuantizedKVCacheEntry | None = None, +) -> torch.Tensor: + """SDPA with optional query scaling. + + Args: + q_BSHD (torch.Tensor): Queries of shape (batch, seq len, num heads, head dim). + k_BSJD (torch.Tensor | None): Keys of shape (batch, seq len, num heads or + num kv heads, head dim). None when `quantized_kv` carries them. + v_BSJD (torch.Tensor | None): Values of shape (batch, seq len, num heads + or num kv heads, head dim). None when `quantized_kv` carries them. + softmax_scaling_layer (nn.Module | None): Optional module to apply + SSMax scaling to queries before attention. `n` for the scaling is the + KV sequence length, read from `k_BSJD` or the quantized cache entry. + _backends_override (list[SDPBackend] | None): Optional list of SDP backends. + quantized_kv (QuantizedKVCacheEntry | None): Keys and values as a + quantized cache entry, in place of `k_BSJD`/`v_BSJD`. The dispatcher + dequantizes it unless the backend it picks consumes it as stored. + + Returns: + torch.Tensor: Attention output, shape (B, S, H, D). + """ + if softmax_scaling_layer is not None: + k = quantized_kv.key if quantized_kv is not None else k_BSJD + assert k is not None + src_len = k.shape[1] + q_BSHD = softmax_scaling_layer(q_BSHD, src_len) + return scaled_dot_product_attention( + q_BSHD, + k_BSJD, + v_BSJD, + _backends_override, + quantized_kv=quantized_kv, + ) + + +# --------------------------------------------------------------------------- +# Attention modules +# --------------------------------------------------------------------------- + + +class Attention(nn.Module): + """Multi-head self-attention with RoPE.""" + + def __init__( + self, + embedding_size: int, + num_heads: int, + head_dim: int, + *, + norm_factory: Callable[[int], nn.Module], + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + kw = {"device": device, "dtype": dtype, "bias": False} + + self.q_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.k_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.v_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.out_projection = nn.Linear(head_dim * num_heads, embedding_size, **kw) + + torch.nn.init.xavier_uniform_(self.q_projection.weight) + torch.nn.init.xavier_uniform_(self.k_projection.weight) + torch.nn.init.xavier_uniform_(self.v_projection.weight) + torch.nn.init.zeros_(self.out_projection.weight) + + self.q_norm = norm_factory(head_dim) + self.k_norm = norm_factory(head_dim) + + @override + def forward(self, x_BSE: torch.Tensor, rope: RotaryEmbedding) -> torch.Tensor: + B, S, _ = x_BSE.shape + q = self.q_projection(x_BSE).view(B, S, -1, self.head_dim) + k = self.k_projection(x_BSE).view(B, S, -1, self.head_dim) + v = self.v_projection(x_BSE).view(B, S, -1, self.head_dim) + + q = rope.rotate_queries_or_keys(q.transpose(1, 2)).transpose(1, 2) + k = rope.rotate_queries_or_keys(k.transpose(1, 2)).transpose(1, 2) + q = self.q_norm(q) + k = self.k_norm(k) + + out = _batched_scaled_dot_product_attention(q, k, v).reshape( + B, S, self.head_dim * self.num_heads + ) + return self.out_projection(out) + + +class CrossAttention(nn.Module): + """Multi-head cross-attention (query attends to key/value sequence).""" + + def __init__( + self, + embedding_size: int, + num_heads: int, + head_dim: int, + softmax_scaling_layer: nn.Module | None = None, + *, + norm_factory: Callable[[int], nn.Module], + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.softmax_scaling_layer = softmax_scaling_layer + kw = {"device": device, "dtype": dtype, "bias": False} + + self.q_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.k_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.v_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.out_projection = nn.Linear(head_dim * num_heads, embedding_size, **kw) + + torch.nn.init.xavier_uniform_(self.q_projection.weight) + torch.nn.init.xavier_uniform_(self.k_projection.weight) + torch.nn.init.xavier_uniform_(self.v_projection.weight) + torch.nn.init.zeros_(self.out_projection.weight) + + self.q_norm = norm_factory(head_dim) + self.k_norm = norm_factory(head_dim) + + @override + def forward( + self, + x_for_query_BQE: torch.Tensor, + x_for_key_and_value_BVE: torch.Tensor, + ) -> torch.Tensor: + B, Q, _ = x_for_query_BQE.shape + _, V, _ = x_for_key_and_value_BVE.shape + q = self.q_projection(x_for_query_BQE).view(B, Q, -1, self.head_dim) + k = self.k_projection(x_for_key_and_value_BVE).view(B, V, -1, self.head_dim) + v = self.v_projection(x_for_key_and_value_BVE).view(B, V, -1, self.head_dim) + + q = self.q_norm(q) + k = self.k_norm(k) + + out = _batched_scaled_dot_product_attention( + q, + k, + v, + softmax_scaling_layer=self.softmax_scaling_layer, + ) + + return self.out_projection(out.reshape(B, Q, self.head_dim * self.num_heads)) + + +class ICLAttention(nn.Module): + """ICL attention: all rows attend to train-only keys/values. + + In v2, the ICL transformer restricts keys/values to training rows so that + test rows cannot attend to each other or to future labels. + + When `num_kv_heads_test` is set, test rows use fewer KV heads than train + rows (GQA / MQA for the test partition only), reducing the KV-cache at + inference time. + """ + + def __init__( + self, + embedding_size: int, + num_heads: int, + head_dim: int, + softmax_scaling_layer: nn.Module | None = None, + num_kv_heads: int | None = None, + num_kv_heads_test: int | None = None, + *, + norm_factory: Callable[[int], nn.Module], + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.softmax_scaling_layer = softmax_scaling_layer + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_heads_test = num_kv_heads_test + kw = {"device": device, "dtype": dtype, "bias": False} + + self.q_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.out_projection = nn.Linear(head_dim * num_heads, embedding_size, **kw) + + torch.nn.init.xavier_uniform_(self.q_projection.weight) + torch.nn.init.zeros_(self.out_projection.weight) + + if num_kv_heads is not None: + # GQA: smaller K/V projections + kv_dim = num_kv_heads * head_dim + self.k_projection = nn.Linear(embedding_size, kv_dim, **kw) + self.v_projection = nn.Linear(embedding_size, kv_dim, **kw) + else: + self.k_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + self.v_projection = nn.Linear(embedding_size, head_dim * num_heads, **kw) + nn.init.xavier_uniform_(self.k_projection.weight) + nn.init.xavier_uniform_(self.v_projection.weight) + + # q/k RMSNorm is applied per head (over head_dim), so it commutes with the + # test-head slicing below; the KV cache stores the already-normed keys. + self.q_norm = norm_factory(head_dim) + self.k_norm = norm_factory(head_dim) + + @override + def forward( + self, + x_BRE: torch.Tensor, + single_eval_pos: int, + *, + cached_kv: KVCacheEntry | QuantizedKVCacheEntry | None = None, + return_kv: bool = False, + ) -> tuple[torch.Tensor, KVCacheEntry | None]: + """Self-attention where k/v are restricted to train rows. + + Args: + x_BRE: (B, R, E) all rows (train + test), or test-only when + `cached_kv` is provided. + single_eval_pos: Number of training rows; positions after this index + are test rows. Should be 0 when using `cached_kv`. + cached_kv: Pre-computed K/V from a previous forward pass. When + provided, K/V projection is skipped and these values are used + directly. + return_kv: If True, also return the computed K/V as a + `KVCacheEntry`. + + Returns: + `(output, kv_entry)` where `kv_entry` is `None` unless + `return_kv` is True. + """ + B, R, _ = x_BRE.shape + + q = self.q_norm( + self.q_projection(x_BRE).view(B, R, self.num_heads, self.head_dim) + ) + + if cached_kv is not None: + # Use pre-computed K/V from cache (test-only path) + k = cached_kv.key + v = cached_kv.value + assert k is not None, "cached key is None" + assert v is not None, "cached value is None" + # The cache already stores only the test KV heads (sliced at + # cache-build time), so no slicing is needed here. + if self.num_kv_heads_test is not None: + nh_test_heads = self.num_kv_heads_test + assert k.shape[2] == nh_test_heads, "cached key has wrong num heads" + assert v.shape[2] == nh_test_heads, "cached value has wrong num heads" + if isinstance(cached_kv, QuantizedKVCacheEntry): + # The SDPA wrapper dequantizes unless a backend takes it as is. + out = _batched_scaled_dot_product_attention( + q, + None, + None, + softmax_scaling_layer=self.softmax_scaling_layer, + quantized_kv=cached_kv, + ) + else: + # Match dtype in case of autocast (e.g. fp32 cache under fp16) + if k.dtype != q.dtype: + k = k.to(q.dtype) + v = v.to(q.dtype) + out = _batched_scaled_dot_product_attention( + q, + k, + v, + softmax_scaling_layer=self.softmax_scaling_layer, + ) + else: + N = R if single_eval_pos is None else single_eval_pos + x_train = x_BRE[:, :N] + k = self.k_projection(x_train).view(B, N, self.num_kv_heads, self.head_dim) + v = self.v_projection(x_train).view(B, N, self.num_kv_heads, self.head_dim) + # Norm once here so the value cached below is already normed. + k = self.k_norm(k) + + if ( + self.num_kv_heads_test is not None + and single_eval_pos is not None + and N < R + ): + # Train rows: full KV heads + out_train = _batched_scaled_dot_product_attention( + q[:, :N], + k, + v, + softmax_scaling_layer=self.softmax_scaling_layer, + ) + # Test rows: fewer KV heads (GQA / MQA) + nh_test_heads = self.num_kv_heads_test + out_test = _batched_scaled_dot_product_attention( + q[:, N:], + k[:, :, :nh_test_heads], + v[:, :, :nh_test_heads], + softmax_scaling_layer=self.softmax_scaling_layer, + ) + out = torch.cat([out_train, out_test], dim=1) + else: + out = _batched_scaled_dot_product_attention( + q, + k, + v, + softmax_scaling_layer=self.softmax_scaling_layer, + ) + + result = self.out_projection(out.reshape(B, R, self.head_dim * self.num_heads)) + + kv_entry: KVCacheEntry | None = None + if return_kv: + # Only cache the KV heads used for test<-train attention to save + # memory. When num_kv_heads_test is set, test rows use fewer heads. + # Under autocast `k_norm` (an RMSNorm) returns fp32 while `v` is the + # autocast dtype; SDPA casts both to that dtype anyway, so store the + # keys at it too rather than at twice the size. `kv_compute_dtype` in + # `forward` is read off the cached key, so this also sets the dtype of + # the decoder keys and of the quantization scales. + k_cache, v_cache = k.to(v.dtype), v + if self.num_kv_heads_test is not None: + nh_test_heads = self.num_kv_heads_test + # .contiguous() so the kept slice owns its storage and the + # full-projection backing tensor can be freed; otherwise the + # cache silently retains all KV heads via the slice view. + k_cache = k_cache[:, :, :nh_test_heads].contiguous() + v_cache = v_cache[:, :, :nh_test_heads].contiguous() + kv_entry = KVCacheEntry(key=k_cache.detach(), value=v_cache.detach()) + return result, kv_entry + + +# --------------------------------------------------------------------------- +# Transformer blocks +# --------------------------------------------------------------------------- + + +class CrossAttentionBlock(nn.Module): + """Cross-attention block with pre-norm and MLP.""" + + def __init__( + self, + *, + emsize: int, + nhead: int, + dim_feedforward: int, + norm_factory: Callable[[int], nn.Module], + softmax_scaling_layer: nn.Module | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + assert emsize % nhead == 0 + kw = {"device": device, "dtype": dtype} + + self.attn = CrossAttention( + embedding_size=emsize, + num_heads=nhead, + head_dim=emsize // nhead, + softmax_scaling_layer=softmax_scaling_layer, + norm_factory=norm_factory, + **kw, + ) + self.mlp = MLP(emsize, dim_feedforward, **kw) + self.layernorm_q = norm_factory(emsize) + self.layernorm_kv = norm_factory(emsize) + self.layernorm2 = norm_factory(emsize) + + @override + def forward( + self, + x_BQE: torch.Tensor, + context_BVE: torch.Tensor, + ) -> torch.Tensor: + attn_out = self.attn( + self.layernorm_q(x_BQE), + self.layernorm_kv(context_BVE), + ) + x_BQE = x_BQE + attn_out + mlp_out = self.mlp(self.layernorm2(x_BQE)) + return x_BQE + mlp_out + + +class TransformerBlock(nn.Module): + """Standard pre-norm transformer block used in ColumnAggregator.""" + + def __init__( + self, + *, + emsize: int, + nhead: int, + dim_feedforward: int, + norm_factory: Callable[[int], nn.Module], + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + kw = {"device": device, "dtype": dtype} + assert emsize % nhead == 0 + self.attention = Attention( + embedding_size=emsize, + num_heads=nhead, + head_dim=emsize // nhead, + norm_factory=norm_factory, + **kw, + ) + self.layernorm = norm_factory(emsize) + self.layernorm_mlp = norm_factory(emsize) + self.mlp = MLP(emsize, dim_feedforward, **kw) + + @override + def forward( + self, + x_BRCE: torch.Tensor, + rope: RotaryEmbedding, + save_peak_memory_factor: int | None = None, + ) -> torch.Tensor: + x_BRCE = chunked_evaluate_maybe_inplace( + lambda x, rope: self.attention(self.layernorm(x), rope=rope), + x_BRCE, + save_peak_memory_factor=save_peak_memory_factor, + residual=True, + batch_dims=2, + rope=rope, + ) + return chunked_evaluate_maybe_inplace( + lambda x: self.mlp(self.layernorm_mlp(x)), + x_BRCE, + save_peak_memory_factor=save_peak_memory_factor, + residual=True, + batch_dims=3, + ) + + def forward_cross( + self, + query_BRQE: torch.Tensor, + context_BRCE: torch.Tensor, + rope: RotaryEmbedding, + ) -> torch.Tensor: + """Cross-attention variant: query attends to context. + + Used in ColumnAggregator for the last CLS-readout block. + """ + B, R, Q, _ = query_BRQE.shape + _, _, V, E = context_BRCE.shape + + # Fold rows into batch for attention (per-row cross-attn over features) + norm_q = self.layernorm(query_BRQE) + q_flat = norm_q.view(B * R, Q, E) + c_flat = self.layernorm(context_BRCE).view(B * R, V, E) + q_proj = self.attention.q_projection(q_flat).view( + B * R, Q, -1, self.attention.head_dim + ) + k_flat = self.attention.k_projection(c_flat).view( + B * R, V, -1, self.attention.head_dim + ) + v_flat = self.attention.v_projection(c_flat).view( + B * R, V, -1, self.attention.head_dim + ) + + q_proj = rope.rotate_queries_or_keys(q_proj.transpose(1, 2)).transpose(1, 2) + k_flat = rope.rotate_queries_or_keys(k_flat.transpose(1, 2)).transpose(1, 2) + q_proj = self.attention.q_norm(q_proj) + k_flat = self.attention.k_norm(k_flat) + + attn_out = _batched_scaled_dot_product_attention(q_proj, k_flat, v_flat) + attn_out = attn_out.reshape( + B * R, Q, self.attention.head_dim * self.attention.num_heads + ) + attn_out = self.attention.out_projection(attn_out).view(B, R, Q, E) + + x_out = query_BRQE + attn_out + mlp_out = self.mlp(self.layernorm_mlp(x_out)) + return x_out + mlp_out + + +class ICLTransformerBlock(nn.Module): + """ICL transformer block with train-only keys and optional softmax scaling.""" + + def __init__( + self, + *, + emsize: int, + nhead: int, + dim_feedforward: int, + norm_factory: Callable[[int], nn.Module], + softmax_scaling_layer: nn.Module | None = None, + num_kv_heads: int | None = None, + num_kv_heads_test: int | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + kw = {"device": device, "dtype": dtype} + assert emsize % nhead == 0 + self.icl_attention = ICLAttention( + embedding_size=emsize, + num_heads=nhead, + head_dim=emsize // nhead, + softmax_scaling_layer=softmax_scaling_layer, + num_kv_heads=num_kv_heads, + num_kv_heads_test=num_kv_heads_test, + norm_factory=norm_factory, + **kw, + ) + self.layernorm = norm_factory(emsize) + self.layernorm_mlp = norm_factory(emsize) + self.mlp = MLP(emsize, dim_feedforward, **kw) + + @override + def forward( + self, + x_BRE: torch.Tensor, + single_eval_pos: int, + save_peak_memory_factor: int | None = None, + *, + cached_kv: KVCacheEntry | QuantizedKVCacheEntry | None = None, + return_kv: bool = False, + ) -> tuple[torch.Tensor, KVCacheEntry | None]: + """Forward pass with optional KV cache support. + + Args: + x_BRE: (B, R, E) all rows, or test-only when `cached_kv` is set. + single_eval_pos: Number of training rows. + save_peak_memory_factor: Chunking factor for memory saving. + cached_kv: Pre-computed K/V for this layer. + return_kv: If True, also return the K/V cache entry. + + Returns: + `(output, kv_entry)` where `kv_entry` is `None` unless + `return_kv` is True. + """ + kv_entry: KVCacheEntry | None = None + + if return_kv: + # Run attention without chunking so we can capture the KV entry + attn_out, kv_entry = self.icl_attention( + self.layernorm(x_BRE), + single_eval_pos=single_eval_pos, + return_kv=True, + ) + x_BRE = x_BRE + attn_out + elif cached_kv is not None: + # Use cached KV -- chunking over test batch is fine + # TODO: Performance test this as it might not be needed. + def _attn_fn_cached( + x: torch.Tensor, + single_eval_pos: int | None = None, + ) -> torch.Tensor: + out, _ = self.icl_attention( + self.layernorm(x), + single_eval_pos=single_eval_pos, + cached_kv=cached_kv, + ) + return out + + x_BRE = chunked_evaluate_maybe_inplace( + _attn_fn_cached, + x_BRE, + save_peak_memory_factor=save_peak_memory_factor, + residual=True, + batch_dims=1, + single_eval_pos=single_eval_pos, + ) + else: + # Default path -- no cache + def _attn_fn( + x: torch.Tensor, + single_eval_pos: int | None = None, + ) -> torch.Tensor: + out, _ = self.icl_attention( + self.layernorm(x), + single_eval_pos=single_eval_pos, + ) + return out + + x_BRE = chunked_evaluate_maybe_inplace( + _attn_fn, + x_BRE, + save_peak_memory_factor=save_peak_memory_factor, + residual=True, + batch_dims=1, + single_eval_pos=single_eval_pos, + ) + + # MLP (always the same regardless of cache mode) + x_BRE = chunked_evaluate_maybe_inplace( + lambda x: self.mlp(self.layernorm_mlp(x)), + x_BRE, + save_peak_memory_factor=save_peak_memory_factor, + residual=True, + batch_dims=2, + ) + + return x_BRE, kv_entry + + +# --------------------------------------------------------------------------- +# Induced self-attention block (v2 style, no affine output) +# --------------------------------------------------------------------------- + + +class InducedSelfAttentionBlock(nn.Module): + """Induced self-attention (SetTransformer-style) for efficient O(n) attention. + + Two-stage mechanism: + 1. Inducing points attend to train rows uses softmax scaling when provided. + 2. All rows attend to the inducing-point hidden states. + """ + + def __init__( + self, + *, + emsize: int, + nhead: int, + num_inducing_points: int, + dim_feedforward: int, + norm_factory: Callable[[int], nn.Module], + softmax_scaling_layer: nn.Module | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + kw = {"device": device, "dtype": dtype} + block_kw = { + "emsize": emsize, + "nhead": nhead, + "dim_feedforward": dim_feedforward, + "norm_factory": norm_factory, + **kw, + } + + self.cross_attn_block1 = CrossAttentionBlock( + **block_kw, + softmax_scaling_layer=softmax_scaling_layer, + ) + self.cross_attn_block2 = CrossAttentionBlock(**block_kw) + + self.num_inducing_points = num_inducing_points + self.inducing_vectors = nn.Parameter(torch.empty(num_inducing_points, emsize)) + nn.init.trunc_normal_(self.inducing_vectors, std=0.02) + + def _induced_attention( + self, + x_BcRE: torch.Tensor, + single_eval_pos: int | None = None, + cached_hidden: torch.Tensor | None = None, + *, + return_hidden: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Induced self-attention with optional hidden-state return. + + When `return_hidden` is True, returns `(output, hidden_detached)` + so the caller can cache the inducing hidden states. Here, we opt for + different output types depending on return_hidden, so that this function + can be used in `chunked_evaluate_maybe_inplace` without any additional logic. + """ + if cached_hidden is not None: + hidden = cached_hidden.to(x_BcRE.dtype) + else: + Bc, R, _ = x_BcRE.shape + N = R if single_eval_pos is None else single_eval_pos + ind = self.inducing_vectors.unsqueeze(0).expand(Bc, -1, -1) + hidden = self.cross_attn_block1(ind, x_BcRE[:, :N]) + out = self.cross_attn_block2(x_BcRE, hidden) + if return_hidden: + return out, hidden.detach() + return out + + @override + def forward( + self, + x_BRCE: torch.Tensor, + single_eval_pos: int | None = None, + save_peak_memory_factor: int | None = None, + *, + cached_hidden: torch.Tensor | None = None, + return_hidden: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward with optional inducing hidden-state caching. + + Returns: + `(output, hidden)` where `hidden` is `None` unless + `return_hidden` is True. + """ + B, R, C, E = x_BRCE.shape + x_BCRE = x_BRCE.transpose(1, 2).contiguous() + x_BcRE = x_BCRE.reshape(B * C, R, E) + + if return_hidden: + out_BcRE, hidden = self._induced_attention( + x_BcRE, + single_eval_pos=single_eval_pos, + return_hidden=True, + ) + else: + out_BcRE = chunked_evaluate_maybe_inplace( + self._induced_attention, + x_BcRE, + save_peak_memory_factor, + residual=False, + batch_dims=1, + single_eval_pos=single_eval_pos, + cached_hidden=cached_hidden, + ) + hidden = None + + out_BCRE = out_BcRE.reshape(B, C, R, E) + return out_BCRE.transpose(1, 2).contiguous(), hidden + + +# --------------------------------------------------------------------------- +# Feature distribution embedder +# --------------------------------------------------------------------------- + + +class FeatureDistributionEmbedder(nn.Module): + """Stack of InducedSelfAttentionBlock layers applied per column.""" + + def __init__( + self, + *, + emsize: int, + nhead: int, + num_inducing_points: int, + dim_feedforward: int, + num_layers: int, + norm_factory: Callable[[int], nn.Module], + softmax_scaling_layer_factory: Callable[[], nn.Module] | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ): + super().__init__() + self.layers = nn.ModuleList( + InducedSelfAttentionBlock( + emsize=emsize, + nhead=nhead, + num_inducing_points=num_inducing_points, + dim_feedforward=dim_feedforward, + norm_factory=norm_factory, + softmax_scaling_layer=( + softmax_scaling_layer_factory() + if softmax_scaling_layer_factory is not None + else None + ), + device=device, + dtype=dtype, + ) + for _ in range(num_layers) + ) + + @override + def forward( + self, + x_BRiCE: torch.Tensor, + num_train_rows: int | None = None, + save_peak_memory_factor: int | None = None, + *, + force_recompute_layer: bool = False, + cached_hidden: list[torch.Tensor] | None = None, + return_hidden: bool = False, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + """Forward pass through all induced self-attention blocks. + + Returns: + `(output, hidden_states)` where `hidden_states` is `None` + unless `return_hidden` is True. + """ + hidden_states: list[torch.Tensor] | None = [] if return_hidden else None + assert not (return_hidden and force_recompute_layer), ( + "return_hidden is incompatible with force_recompute_layer" + ) + for i, layer in enumerate(self.layers): + if force_recompute_layer: + x_BRiCE, _ = torch.utils.checkpoint.checkpoint( # type: ignore + layer, + x_BRiCE, + num_train_rows, + use_reentrant=False, + save_peak_memory_factor=save_peak_memory_factor, + ) + else: + layer_cached = cached_hidden[i] if cached_hidden is not None else None + x_BRiCE, h = layer( + x_BRiCE, + single_eval_pos=num_train_rows, + save_peak_memory_factor=save_peak_memory_factor, + cached_hidden=layer_cached, + return_hidden=return_hidden, + ) + if hidden_states is not None: + hidden_states.append(h) + return x_BRiCE, hidden_states + + +# --------------------------------------------------------------------------- +# Cross-feature interaction (Row interaction / v2 RowInteraction) +# --------------------------------------------------------------------------- + + +class ColumnAggregator(nn.Module): + """Context-aware cross-feature interaction that aggregates column information. + + CLS tokens are prepended, the sequence passes through transformer blocks, + and the last block performs CLS-only readout (q=CLS, k/v=all). + An output normalization is applied before the CLS tokens are returned. + """ + + def __init__( + self, + emsize: int, + nhead: int, + num_layers: int, + dim_feedforward: int, + num_cls_tokens: int, + *, + norm_factory: Callable[[int], nn.Module], + rope_base: float = 100_000, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + self.embed_dim = emsize + self.num_cls_tokens = num_cls_tokens + kw = {"device": device, "dtype": dtype} + + self.blocks = nn.ModuleList( + TransformerBlock( + emsize=emsize, + nhead=nhead, + dim_feedforward=dim_feedforward, + norm_factory=norm_factory, + **kw, + ) + for _ in range(num_layers) + ) + self.rope = RotaryEmbedding( + dim=emsize // nhead, theta=int(rope_base), interleaved=False + ) + self.cls_tokens = nn.Parameter(torch.empty(num_cls_tokens, emsize)) + nn.init.trunc_normal_(self.cls_tokens, std=0.02) + + # Output norm applied to CLS tokens after the last block (v2 out_ln) + self.out_ln = norm_factory(emsize) + + @override + def forward( + self, + x_BRiCE: torch.Tensor, + save_peak_memory_factor: int | None = None, + force_recompute_layer: bool = False, + ) -> torch.Tensor: + """Transform feature embeddings into per-row CLS representations. + + Args: + x_BRiCE: (B, Ri, C, E) + save_peak_memory_factor: If set, chunk the evaluation to save memory. + force_recompute_layer: If True, force gradient checkpointing. + + Returns: + (B, Ri, num_cls_tokens, E) + """ + B, Ri, _, E = x_BRiCE.shape + cls = self.cls_tokens.expand(B, Ri, self.num_cls_tokens, E).to(x_BRiCE.device) + # Prepend CLS tokens: (B, Ri, num_cls + C, E) + x = torch.cat((cls, x_BRiCE), dim=2) + + # Run all blocks except the last + for block in self.blocks[:-1]: + if force_recompute_layer: + x = torch.utils.checkpoint.checkpoint( # type: ignore + block, + x, + self.rope, + save_peak_memory_factor, + use_reentrant=False, + ) + else: + x = block( + x, rope=self.rope, save_peak_memory_factor=save_peak_memory_factor + ) + + # Last block: CLS tokens as query, full sequence as key/value (v2 readout) + last_block = cast("TransformerBlock", self.blocks[-1]) + x_full: torch.Tensor = x # type: ignore[assignment] + cls_part = x_full[..., : self.num_cls_tokens, :] + if force_recompute_layer: + cls_out = torch.utils.checkpoint.checkpoint( # type: ignore + last_block.forward_cross, + cls_part, + x_full, + self.rope, + use_reentrant=False, + ) + else: + cls_out = last_block.forward_cross(cls_part, x_full, self.rope) + + del x + return self.out_ln(cls_out) + + +class _PreHeadMLP(nn.Module): + """Residual pre-norm MLP applied to embeddings before an output head. + + At init the inner MLP's final projection is zero, so this block acts as + the identity and existing trained heads see unchanged embeddings. + """ + + def __init__( + self, + emsize: int, + dim_feedforward: int, + norm_factory: Callable[[int], nn.Module], + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ) -> None: + super().__init__() + self.norm = norm_factory(emsize) + self.mlp = MLP(emsize, dim_feedforward, device=device, dtype=dtype) + + @override + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.mlp(self.norm(x)) + + +class MultiTaskHeads(nn.Module): + """Bundled output heads for multitask inference. + + Two heads share an input embedding of size `input_size`: + - multiclass: attention-based `ManyClassDecoder` over the train targets, + producing `(M, B, max_num_classes)` logits. + - regression: linear layer producing `(M, B, num_buckets)` bar logits. + + A residual pre-norm MLP block is applied to the embeddings before each final + projection — one for multiclass, a separate one for regression. + """ + + def __init__( + self, + *, + input_size: int, + max_num_classes: int, + num_buckets: int, + decoder_head_dim: int, + decoder_num_heads: int, + decoder_softmax_scaling_layer: nn.Module | None = None, + mlp_dim_feedforward: int, + norm_factory: Callable[[int], nn.Module], + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ): + super().__init__() + kw = {"device": device, "dtype": dtype} + self.many_class_decoder = ManyClassDecoder( + max_num_classes=max_num_classes, + input_size=input_size, + head_dim=decoder_head_dim, + num_heads=decoder_num_heads, + softmax_scaling_layer=decoder_softmax_scaling_layer, + ) + self.output_projection = nn.Linear(input_size, num_buckets, **kw) + self.mlp_classification = _PreHeadMLP( + emsize=input_size, + dim_feedforward=mlp_dim_feedforward, + norm_factory=norm_factory, + **kw, + ) + self.mlp_regression = _PreHeadMLP( + emsize=input_size, + dim_feedforward=mlp_dim_feedforward, + norm_factory=norm_factory, + **kw, + ) + self.register_buffer( + "regression_borders", + _spline_based_regression_borders(num_buckets), + ) + + def project_decoder_keys(self, train_emb: torch.Tensor) -> torch.Tensor: + """Many-class decoder keys for the train rows, ready to cache. + + Runs the classification pre-head MLP first, so the keys come from the same + embeddings `forward` would have projected. + """ + return self.many_class_decoder.project_keys(self.mlp_classification(train_emb)) + + @override + def forward( + self, + train_keys_BNHD: torch.Tensor | None, # from project_decoder_keys + test_emb: torch.Tensor, # (B, M, D) + y_train_BN: torch.Tensor, # (B, N), only consumed by the multiclass head + *, + task_type: str, + num_present_classes: int | None, + ) -> torch.Tensor: + """Apply the head selected by `task_type`. + + Returns `(M, B, max_num_classes)` for multiclass and `(M, B, + num_buckets)` for regression. `train_keys_BNHD` is unused for regression, + which has no many-class decoder, and may be None there. + """ + if task_type == "regression": + test_emb = self.mlp_regression(test_emb) + return self.output_projection(test_emb.transpose(0, 1)) + if task_type == "multiclass": + assert num_present_classes is not None + assert train_keys_BNHD is not None, ( + "the multiclass head needs the decoder keys" + ) + test_emb = self.mlp_classification(test_emb) + return self.many_class_decoder( + train_keys_BNHD, + test_emb, + y_train_BN, + num_present_classes=num_present_classes, + ) + raise ValueError(f"Unsupported task type: {task_type}") + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + + +class TabPFNV3p5(Architecture): + """Single-file TabPFN v3.5 architecture. + + Pipeline: + 1. Preprocessing: standard scaling + NaN encoding + 2. Feature grouping: circular shifts applied before embedding + 3. Cell embedding: feature_group_size scalar values → embed_dim + 4. Target-aware column embedding: add y_encoder(y_train) to train rows + 5. Feature distribution embedder: InducedSelfAttentionBlock x dist_embed_num_blocks + 6. Feature aggregator with feature interaction: transformer with aggregation tokens + 7. ICL transformer: y_encoder + standard attention (train-keys only) + decoder + """ + + def __init__( + self, + *, + config: TabPFNV3p5Config, + device: torch.device | str | None = None, + dtype: torch.dtype | str | None = None, + ): + super().__init__() + self.ff_factor = config.ff_factor + self.icl_emsize = config.embed_dim * config.feat_agg_num_cls_tokens + self.max_num_classes = config.max_num_classes + self.feature_group_size = config.feature_group_size + self.ecdf_num_buckets = config.cell_ecdf_num_buckets + kw = {"device": device, "dtype": dtype} + + norm_factory = partial(_DtypeMatchingRMSNorm, device=device, dtype=dtype) + + # ---- Cell embedding (ordinal: grouped raw values → E) ---- + self.x_embed = FourierPlusMetadataFeatureGroupEmbedder( + config.feature_group_size, + config.embed_dim, + config.fourier_encoding_num_frequencies, + ecdf_num_frequencies=config.cell_ecdf_num_frequencies, + row_chunk_size=config.cell_embed_row_chunk_size, + **kw, + ) + + # ---- Target-aware col embedding (one per task type) ---- + self.col_y_encoder = nn.ModuleDict( + { + "multiclass": TrainableOrthogonalEmbedding( + config.max_num_classes, + config.embed_dim, + ), + "regression": nn.Linear(1, config.embed_dim, bias=True, **kw), + } + ) + # Shared across task types so that the encoded train-label signal has + # comparable magnitude regardless of which task-specific y-encoder + # produced it. + self.col_y_layernorm = nn.LayerNorm(config.embed_dim, **kw) + + # ---- Distribution embedder (SetTransformer per feature column) ---- + self.feature_distribution_embedder = FeatureDistributionEmbedder( + emsize=config.embed_dim, + nhead=config.dist_embed_num_heads, + num_layers=config.dist_embed_num_blocks, + num_inducing_points=config.dist_embed_num_inducing_points, + dim_feedforward=config.embed_dim * config.ff_factor, + norm_factory=norm_factory, + softmax_scaling_layer_factory=lambda: SoftmaxScalingMLP( + num_heads=config.dist_embed_num_heads, + head_dim=config.embed_dim // config.dist_embed_num_heads, + n_hidden=config.softmax_scaling_mlp_hidden_dim, + ), + **kw, + ) + + # ---- Cross-feature interaction (RowInteraction) ---- + self.column_aggregator = ColumnAggregator( + emsize=config.embed_dim, + nhead=config.feat_agg_num_heads, + num_layers=config.feat_agg_num_blocks, + num_cls_tokens=config.feat_agg_num_cls_tokens, + dim_feedforward=config.embed_dim * config.ff_factor, + norm_factory=norm_factory, + rope_base=config.feat_agg_rope_base, + **kw, + ) + + # ---- ICL target encoder (one per task type) ---- + self.icl_y_encoder = nn.ModuleDict( + { + "multiclass": TrainableOrthogonalEmbedding( + config.max_num_classes, + self.icl_emsize, + ), + "regression": nn.Linear(1, self.icl_emsize, bias=True, **kw), + } + ) + self.icl_y_layernorm = nn.LayerNorm(self.icl_emsize, **kw) + + # ---- ICL transformer ---- + self.icl_blocks = nn.ModuleList( + ICLTransformerBlock( + emsize=self.icl_emsize, + nhead=config.icl_num_heads, + dim_feedforward=self.icl_emsize * config.ff_factor, + norm_factory=norm_factory, + num_kv_heads=config.icl_num_kv_heads, + num_kv_heads_test=config.icl_num_kv_heads_test, + softmax_scaling_layer=SoftmaxScalingMLP( + num_heads=config.icl_num_heads, + head_dim=self.icl_emsize // config.icl_num_heads, + n_hidden=config.softmax_scaling_mlp_hidden_dim, + ), + **kw, + ) + for _ in range(config.nlayers) + ) + + # ---- Output norm + multi-task heads ---- + self.output_norm = norm_factory(self.icl_emsize) + decoder_softmax_scaling = ( + SoftmaxScalingMLP( + num_heads=config.decoder_num_heads, + head_dim=config.decoder_head_dim, + n_hidden=config.softmax_scaling_mlp_hidden_dim, + ) + if config.decoder_use_softmax_scaling + else None + ) + self.heads = MultiTaskHeads( + input_size=self.icl_emsize, + max_num_classes=config.max_num_classes, + num_buckets=config.num_buckets, + decoder_head_dim=config.decoder_head_dim, + decoder_num_heads=config.decoder_num_heads, + decoder_softmax_scaling_layer=decoder_softmax_scaling, + mlp_dim_feedforward=self.icl_emsize * config.ff_factor, + norm_factory=norm_factory, + device=device, + dtype=dtype, + ) + # Expose for API compatibility. + self.regression_borders = self.heads.regression_borders + self.standard_scaler = TorchStandardScaler() + self._nan_safe_output = True + self._icl_bf16 = False + self.emsize = config.embed_dim + self.inference_row_chunk_size = config.inference_row_chunk_size + self.inference_col_chunk_size = config.inference_col_chunk_size + + def enable_icl_bf16(self) -> None: + """Switch the ICL blocks and output norm to bfloat16 inference.""" + self.icl_blocks.to(torch.bfloat16) + self.output_norm.to(torch.bfloat16) + self._icl_bf16 = True + + @property + @override + def embedding_dim(self) -> int: + return self.icl_emsize + + @override + def forward( + self, + x: torch.Tensor | dict[str, torch.Tensor], + y: torch.Tensor | dict[str, torch.Tensor] | None, + task_type: TaskType, + *, + only_return_standard_out: bool = True, + categorical_inds: list[list[int]] | None = None, + performance_options: PerformanceOptions | None = None, + kv_cache: TabPFNV3p5Cache | None = None, + return_kv_cache: bool = False, + x_is_test_only: bool = False, + ) -> ( + torch.Tensor + | dict[str, torch.Tensor] + | tuple[torch.Tensor | dict[str, torch.Tensor], TabPFNV3p5Cache | None] + ): + """Main forward pass for TabPFN v3.5. + + `task_type` selects the per-task target encoder and output head; the + same model instance handles both tasks. + + When a KV cache is provided, `x_is_test_only=True` lets the + caller pass only the test rows (shape `(num_test, 1, D)`) instead + of padding with train-row placeholders. `y` still carries the + train labels — the decoder reads `y[:num_train]` for the + many-class head. Outside the cache path, `x` is always the full + dataset and this flag is ignored. + """ + del categorical_inds + if isinstance(x, dict): + x = x["main"] + if isinstance(y, dict): + y = y["main"] + if y is None: + y = torch.zeros(0, device=x.device, dtype=x.dtype) + if y.dim() == 3 and y.shape[-1] == 1: + y = y.squeeze(-1) + + if performance_options is None: + performance_options = self.get_default_performance_options() + + if performance_options.enable_torch_compile: + # We increase the limit, since we compile a couple of subgraphs for + # chunking and different batched_sdpa configs. + torch._dynamo.config.cache_size_limit = max( + 32, torch._dynamo.config.cache_size_limit + ) + + if x_is_test_only and (kv_cache is None or kv_cache.is_empty()): + raise ValueError( + "x_is_test_only=True requires kv_cache to be provided; " + "the non-cache forward needs the full train+test tensor." + ) + + num_present_classes = None + if task_type == "multiclass": + num_present_classes = ( + torch.nan_to_num(y, nan=0.0).max().item() + 1 if y.numel() else 1 + ) + if not self.training and ( + num_present_classes > self.max_num_classes or (y < 0).any() + ): + raise ValueError( + "Target is out of range. " + "Make sure to use an ordinal encoded target. " + f"Expected target values between 0 and {self.max_num_classes - 1}, " + f"but got values outside this range." + ) + num_present_classes = int(num_present_classes) + x_RiBC = x + B = x_RiBC.shape[1] + num_train = y.shape[0] + if performance_options.enable_torch_compile: + torch._dynamo.mark_dynamic(x_RiBC, index=0) + torch._dynamo.mark_dynamic(x_RiBC, index=1) + torch._dynamo.mark_dynamic(x_RiBC, index=2) + + x_BRiClE, inducing_hidden, scaler_stats = self._stages_0_to_2( + x_RiBC, + y, + task_type, + performance_options=performance_options, + return_inducing_hidden=return_kv_cache, + kv_cache=kv_cache, + x_is_test_only=x_is_test_only, + ) + + # ---- Stage 3: ICL ---- + x_BRiD = x_BRiClE.flatten(-2) + del x_BRiClE + if self._icl_bf16: + x_BRiD = x_BRiD.to(torch.bfloat16) + + # Per-layer KV entries collected when return_kv_cache is True. + kv_out: dict[int, KVCacheEntry | QuantizedKVCacheEntry] = {} + # The compute dtype of the K/V, captured before any quantization below. + # The decoder keys are stored at this dtype, not at the quantized one. + kv_compute_dtype: torch.dtype | None = None + # tabpfn releases up to 8.3.0 have no `kv_cache_dtype`; their engine + # quantizes the finished cache through `TabPFNV3p5Cache.quantize`. Newer + # ones set this instead and expect the quantization per layer, here, + # which frees each full-precision entry as the loop moves on. + kv_cache_dtype = getattr(performance_options, "kv_cache_dtype", None) + + # An ambient autocast context (e.g. the TabPFN inference engine wraps + # forward in one) would recast matmuls to the autocast dtype and upcast + # layer_norm to fp32; bf16 ICL needs true bf16 compute throughout. + icl_autocast_ctx = ( + torch.autocast(x_BRiD.device.type, enabled=False) + if self._icl_bf16 + else contextlib.nullcontext() + ) + with icl_autocast_ctx: + if kv_cache is not None and not kv_cache.is_empty(): + # Cache path: no y_icl embedding; use cached K/V pairs + for layer_idx, block in enumerate(self.icl_blocks): + x_BRiD, _ = block( + x_BRiD, + 0, + performance_options.save_peak_memory_factor, + cached_kv=kv_cache.kv[layer_idx], + ) + else: + if num_train > 0: + y_icl = self._prepare_y(y, num_train, B, task_type=task_type) + y_icl_emb = self._embed_icl_y(y_icl, task_type=task_type) + x_BRiD[:, :num_train] = x_BRiD[:, :num_train] + y_icl_emb + + if return_kv_cache: + for layer_idx, block in enumerate(self.icl_blocks): + x_BRiD, kv_entry = block( + x_BRiD, + num_train, + performance_options.save_peak_memory_factor, + return_kv=True, + ) + kv_compute_dtype = kv_entry.key.dtype + if kv_cache_dtype is not None: + kv_entry = kv_entry.quantize(kv_cache_dtype) + kv_out[layer_idx] = kv_entry + else: + for block in self.icl_blocks: + if performance_options.force_recompute_layer: + x_BRiD, _ = torch.utils.checkpoint.checkpoint( + block, + x_BRiD, + num_train, + use_reentrant=False, + save_peak_memory_factor=performance_options.save_peak_memory_factor, + ) + else: + x_BRiD, _ = block( + x_BRiD, + num_train, + performance_options.save_peak_memory_factor, + ) + + x_BRiD = self.output_norm(x_BRiD) + if self._icl_bf16: + assert x_BRiD.dtype == torch.bfloat16, ( + "bf16 ICL inference must preserve a bf16 residual through the " + "ICL blocks and output norm" + ) + # Preserve the fp32 boundary expected by task heads that were not + # moved to bf16. + x_BRiD = x_BRiD.float() + + # ---- Split embeddings -------------------------------------------------- + running_from_cache = kv_cache is not None and not kv_cache.is_empty() + if running_from_cache: + test_emb = x_BRiD + # The cache holds the decoder keys projected from the train + # embeddings, not the embeddings, so they are gone on this path. + train_emb = None + else: + test_emb = x_BRiD[:, num_train:] + train_emb = x_BRiD[:, :num_train] + + # ---- Many-class decoder keys ------------------------------------------- + # Regression has no many-class decoder, so it neither builds nor caches + # these; a regression cache is that much smaller. + train_keys: torch.Tensor | None = None + if task_type == "multiclass": + if running_from_cache: + assert kv_cache is not None + assert kv_cache.decoder_keys is not None, ( + "a multiclass KV cache must carry the decoder keys" + ) + train_keys = kv_cache.decoder_keys + else: + assert train_emb is not None + train_keys = self.heads.project_decoder_keys(train_emb) + + # ---- Build KV cache output --------------------------------------------- + built_cache: TabPFNV3p5Cache | None = None + if return_kv_cache: + if running_from_cache: + built_cache = kv_cache # pass through unchanged + else: + # Reuse the statistics fitted during preprocessing (on the imputed + # train rows). Re-fitting on raw `x_RiBC` here would let the + # passed-through +/-inf poison the mean/std and turn every + # standardised test cell into NaN at predict time. + assert kv_out + # Store the decoder keys at the unquantized ICL compute dtype: + # the KV entries above may already be int8 by now. + assert kv_compute_dtype is not None + built_cache = TabPFNV3p5Cache( + kv=kv_out, + decoder_keys=( + train_keys.detach().to(kv_compute_dtype) + if train_keys is not None + else None + ), + train_shape=(B, num_train), + scaler_cache={ + k: v.detach() + for k, v in scaler_stats.items() + if k != _ECDF_CONTEXT_KEY + }, + ecdf_context=scaler_stats[_ECDF_CONTEXT_KEY].detach(), + inducing_hidden=( + [h.detach() for h in inducing_hidden] + if inducing_hidden is not None + else None + ), + ) + + # ---- Decoder ----------------------------------------------------------- + y_BN = y.transpose(0, 1) if y.dim() == 2 else y.unsqueeze(0) + y_train_BN = y_BN[:, :num_train] + test_out: torch.Tensor = self.heads( + train_keys, + test_emb, + y_train_BN, + task_type=task_type, + num_present_classes=num_present_classes, + ) + if self._nan_safe_output: + test_out = torch.nan_to_num(test_out, nan=0.0) + + if only_return_standard_out: + output = test_out + else: + output = { + "standard": test_out, + "test_embeddings": test_emb.transpose(0, 1), + } + if train_emb is not None: + output["train_embeddings"] = train_emb.transpose(0, 1) + + if return_kv_cache: + return output, built_cache + return output + + @override + def get_default_performance_options(self) -> PerformanceOptions: + options = super().get_default_performance_options() + return dataclasses.replace( + options, + use_chunkwise_inference=True, + ) + + @override + def get_supported_kv_cache_precisions(self) -> tuple[str, ...]: + # `TabPFNV3p5Cache.quantize` handles both dtypes. Without this override the + # base returns ("auto",) and the engine never quantizes. + return ("auto", "int8", "fp8") + + def _prepare_y( + self, + y: torch.Tensor, + num_train: int, + batch_size: int, + *, + task_type: TaskType, + ) -> torch.Tensor: + """Prepare y_train for either target-embedding stage. + + Returns: + Clean y_train of shape (B, train_size), or None if no train rows. + """ + if num_train == 0: + raise ValueError("No training rows available for target embedding.") + + y_NB1 = _prepare_targets(y, num_train, batch_size)[:num_train] + y_NB1 = _impute_target_nan_and_inf( + y_NB1=y_NB1, + task_type=task_type, + num_train_rows=num_train, + ) + return y_NB1.squeeze(-1).transpose(0, 1) # (B, train_size) + + def _embed_col_y(self, y_BN: torch.Tensor, *, task_type: TaskType) -> torch.Tensor: + """Embed y_train for the col stage → (B, T, E).""" + if task_type == "multiclass": + y_emb = self.col_y_encoder["multiclass"](y_BN) + elif task_type == "regression": + y_emb = self.col_y_encoder["regression"](y_BN.unsqueeze(-1)) + else: + raise ValueError(f"Unsupported task type: {task_type}") + return self.col_y_layernorm(y_emb) + + def _embed_icl_y(self, y_BN: torch.Tensor, *, task_type: TaskType) -> torch.Tensor: + """Embed y_train for the ICL stage → (B, T, D).""" + if task_type == "multiclass": + y_emb = self.icl_y_encoder["multiclass"](y_BN) + elif task_type == "regression": + y_emb = self.icl_y_encoder["regression"](y_BN.unsqueeze(-1)) + else: + raise ValueError(f"Unsupported task type: {task_type}") + return self.icl_y_layernorm(y_emb) + + def _preprocess_raw( + self, + x_RiBC: torch.Tensor, + num_train: int, + scaler_cache: dict[str, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: + """NaN indicator capture → imputation → standardisation → transpose. + + When *scaler_cache* is provided the scaler is applied without refitting + (inference mode); otherwise it is fitted on the first *num_train* rows + *after* imputation, so the statistics stay finite even when the raw input + carried +/-inf. + + `ecdf_BRiC` holds the per-cell midrank ECDF against the train rows, one + raw rank per cell — the cell embedder lifts it to sin/cos features. + + Returns `(x_BRiC, nan_ind_BRiC, ecdf_BRiC, scaler_stats)`. Returning the + fitted statistics lets the caller store exactly these in the inference + cache, so test rows are standardised (and ECDF-ranked) against the same + train context. + """ + # Note: Indicators need to be computed before imputation. + nan_ind_BRiC = _generate_nan_and_inf_indicator(x_RiBC).transpose(0, 1) + + x_RiBC, is_finite_RiBC = _impute_nan_and_inf_with_mean( + x_RiBC, num_train, scaler_cache + ) + fit_stats = scaler_cache is None + if fit_stats: + fit_data = x_RiBC[:num_train] if num_train > 0 else x_RiBC + scaler_cache = self.standard_scaler.fit(fit_data) + # Align the fill value between train and cached test rows: the cached + # path fills from `mean`, which differs from the nanmean by rounding, + # enough to move a filled test cell out of the ECDF tie block. + x_RiBC = torch.where( + is_finite_RiBC, + x_RiBC, + scaler_cache["mean"].unsqueeze(0).expand_as(x_RiBC), + ) + + # Rank the imputed values, not the output of `standard_scaler.transform` + # below: its +/-100 clip would collapse extreme outliers into ties. + x_imputed_BRiC = x_RiBC.transpose(0, 1) + if fit_stats: + scaler_cache[_ECDF_CONTEXT_KEY] = _build_ecdf_context( + x_imputed_BRiC, num_train, self.ecdf_num_buckets + ) + ecdf_BRiC = _in_context_ecdf(x_imputed_BRiC, scaler_cache[_ECDF_CONTEXT_KEY]) + + x_RiBC = self.standard_scaler.transform(x_RiBC, fitted_cache=scaler_cache) + x_BRiC = x_RiBC.transpose(0, 1) + + return x_BRiC, nan_ind_BRiC, ecdf_BRiC, scaler_cache + + def _group_features( + self, + x_BRiC: torch.Tensor, + nan_ind_BRiC: torch.Tensor, + ecdf_BRiC: torch.Tensor, + ) -> torch.Tensor: + """Build the full grouped + indicator-concatenated tensor. + + Layout: standard-scaled values, then the NaN indicators, then the raw ECDF + ranks. The cell embedder slices the values off the front and the ranks off + the back, so those two blocks must stay at their ends. + """ + shifts = [-(2**i) for i in range(self.feature_group_size)] + return torch.cat( + [ + torch.stack([torch.roll(t, shifts=s, dims=2) for s in shifts], dim=-1) + for t in (x_BRiC, nan_ind_BRiC, ecdf_BRiC) + ], + dim=-1, + ) + + def _group_feature_cols( + self, + x_BRiC: torch.Tensor, + nan_ind_BRiC: torch.Tensor, + ecdf_BRiC: torch.Tensor, + col_start: int, + col_end: int, + ) -> torch.Tensor: + """Grouped features for columns `[col_start, col_end)` only. + + Equivalent to `_group_features(x, ind, ecdf)[:, :, col_start:col_end]` — + `torch.roll(x, -s, dims=2)[:, :, c] == x[:, :, (c + s) % C]` — without + materializing the full `(B, Ri, C, G)` tensor. + """ + C = x_BRiC.shape[2] + cols = torch.arange(col_start, col_end, device=x_BRiC.device) + size = self.feature_group_size + idx = [(cols + 2**i) % C for i in range(size)] + return torch.cat( + [ + torch.stack([t[:, :, i] for i in idx], dim=-1) + for t in (x_BRiC, nan_ind_BRiC, ecdf_BRiC) + ], + dim=-1, + ) + + def _compute_all_inducing_hidden( + self, + dist_embedder_layers: nn.ModuleList, + x_train_BNC: torch.Tensor, + nan_ind_train_BNC: torch.Tensor, + ecdf_train_BNC: torch.Tensor, + y_col_emb_BNE: torch.Tensor | None, + col_chunk_size: int, + *, + enable_torch_compile: bool, + ) -> list[torch.Tensor]: + """Pre-compute inducing hidden states for every dist-embedder block. + + Processes columns in chunks of *col_chunk_size* to avoid + materialising `(B*C_out, N_train, embedding_size)` all at once. + Takes the ungrouped train rows and groups each column chunk on the + fly, so the full grouped tensor is never resident. + + Returns one `(B*C, num_inducing, embedding_size)` tensor per block. + """ + num_columns = x_train_BNC.shape[2] + num_train = x_train_BNC.shape[1] + num_blocks = len(dist_embedder_layers) + # I: num inducing vectors. + # Collect (B, Cj, I, E) per column-chunk, per block + hidden_per_block: list[list[torch.Tensor]] = [[] for _ in range(num_blocks)] + + process_col_fn = ( + self._compiled(self._process_col_chunk) + if enable_torch_compile + else self._process_col_chunk + ) + + for c0 in range(0, num_columns, col_chunk_size): + c1 = min(c0 + col_chunk_size, num_columns) + x_grouped_chunk_BNCjG = self._group_feature_cols( + x_train_BNC, nan_ind_train_BNC, ecdf_train_BNC, c0, c1 + ) + if enable_torch_compile: + torch._dynamo.mark_dynamic(x_grouped_chunk_BNCjG, index=0) + torch._dynamo.mark_dynamic(x_grouped_chunk_BNCjG, index=1) + # Will compile two versions: one with cols dynamic and one with + # cols static for the fixed chunk size. + if (c1 - c0) != col_chunk_size: + torch._dynamo.mark_dynamic(x_grouped_chunk_BNCjG, index=2) + + chunk_outputs_BCjIE = process_col_fn( + x_grouped_chunk_BNCjG=x_grouped_chunk_BNCjG, + y_col_emb_BNE=y_col_emb_BNE, + num_train=num_train, + ) + for blk_idx, h in enumerate(chunk_outputs_BCjIE): + hidden_per_block[blk_idx].append(h) + + # Concatenate and flatten column chunks (B * C_out, I, E) per block. + return [torch.cat(chunks, dim=1).flatten(0, 1) for chunks in hidden_per_block] + + def _compiled(self, method: Callable) -> Callable: + """Lazily `torch.compile` a bound method of this instance. + + The compiled callable is cached per underlying function, so dynamo / + inductor are only imported when `torch.compile` is actually + requested (keeping `import tabpfn` and eager inference free of them), + and each method is compiled at most once. + """ + cache = self.__dict__.setdefault("_torch_compile_cache", {}) + key = method.__func__ + if key not in cache: + cache[key] = torch.compile(method, dynamic=True) + return cache[key] + + def __getstate__(self) -> dict[str, Any]: + # `torch.compile`-d callables are not picklable, so exclude the lazily + # populated compile cache from (un)pickling / torch.save. It is + # rebuilt on demand by `_compiled()`. Delegate to nn.Module first so + # its own state handling (e.g. `_compiled_call_impl`) is preserved. + state = super().__getstate__() + state.pop("_torch_compile_cache", None) + return state + + def _preprocess_and_group( + self, + rows_RiBC: torch.Tensor, + y: torch.Tensor, + num_train: int, + scaler_cache: dict[str, torch.Tensor] | None, + task_type: TaskType, + ) -> tuple[torch.Tensor, torch.Tensor | None, dict[str, torch.Tensor]]: + """Preprocess rows, embed y for the col stage, and group features. + + Combines the three pre-chunk-loop steps into one compiled pass. + Returns the grouped x of shape `(B, Ri, C, G)` tensor, optionally the + `(B, N_train, E)` y embedding, and the scaler statistics fitted during + preprocessing (for reuse in the inference cache). + """ + x_BRiC, nan_ind_BRiC, ecdf_BRiC, y_col_emb_BNE, scaler_stats = ( + self._preprocess_no_group(rows_RiBC, y, num_train, scaler_cache, task_type) + ) + x_grouped_BRiCG = self._group_features(x_BRiC, nan_ind_BRiC, ecdf_BRiC) + return x_grouped_BRiCG, y_col_emb_BNE, scaler_stats + + def _preprocess_no_group( + self, + rows_RiBC: torch.Tensor, + y: torch.Tensor, + num_train: int, + scaler_cache: dict[str, torch.Tensor] | None, + task_type: TaskType, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + dict[str, torch.Tensor], + ]: + """Preprocess rows and embed y, deferring feature grouping. + + Used on the chunked path so the `(B, Ri, C, G)` grouped tensor is never + materialized for all rows at once — chunks are grouped on the fly, + keeping only the `(B, Ri, C)` scaled features, NaN indicators and ECDF + ranks resident. + """ + B = rows_RiBC.shape[1] + x_BRiC, nan_ind_BRiC, ecdf_BRiC, scaler_stats = self._preprocess_raw( + rows_RiBC, num_train, scaler_cache + ) + y_col_emb_BNE: torch.Tensor | None = None + if scaler_cache is None and num_train > 0: + y_col_BN = self._prepare_y(y, num_train, B, task_type=task_type) + y_col_emb_BNE = self._embed_col_y(y_col_BN, task_type=task_type) + return x_BRiC, nan_ind_BRiC, ecdf_BRiC, y_col_emb_BNE, scaler_stats + + def _stages_0_to_2( + self, + x_RiBC: torch.Tensor, + y: torch.Tensor, + task_type: TaskType, + *, + performance_options: PerformanceOptions, + return_inducing_hidden: bool, + kv_cache: TabPFNV3p5Cache | None, + x_is_test_only: bool, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None, dict[str, torch.Tensor]]: + """Stages 0-2: feature embedding, distribution embedding, column aggregation. + + Handles all three computation paths (cache / chunked / full) and returns + `(x_BRiClE, inducing_hidden, scaler_stats)`. `inducing_hidden` is + `None` unless + `return_inducing_hidden` is True (full path) or row-chunking is active + (chunked path, where it is always computed as an intermediate). + """ + num_train = y.shape[0] + if performance_options.use_chunkwise_inference and not self.training: + row_chunk_size = self.inference_row_chunk_size + col_chunk_size = self.inference_col_chunk_size + else: + row_chunk_size = None + col_chunk_size = None + + force_recompute_layer = performance_options.force_recompute_layer + save_peak_memory_factor = performance_options.save_peak_memory_factor + + if kv_cache is not None and not kv_cache.is_empty(): + rows_RiBC = x_RiBC if x_is_test_only else x_RiBC[num_train:] + assert kv_cache.scaler_cache is not None + assert kv_cache.ecdf_context is not None + scaler_cache = { + **kv_cache.scaler_cache, + _ECDF_CONTEXT_KEY: kv_cache.ecdf_context, + } + precomputed_hidden: list[torch.Tensor] | None = kv_cache.inducing_hidden + effective_num_train = 0 + else: + rows_RiBC = x_RiBC + scaler_cache = None + precomputed_hidden = None + effective_num_train = num_train + + num_rows, C = rows_RiBC.shape[0], rows_RiBC.shape[2] + use_chunks = row_chunk_size is not None and row_chunk_size < num_rows + + # --- Preprocess + y col-embed (+ feature grouping on the full path). --- + # The chunked path defers grouping to the per-chunk loops so the full + # (B, Ri, C, G) tensor is never resident. + x_grouped_BRiCG: torch.Tensor | None = None + x_BRiC: torch.Tensor | None = None + nan_ind_BRiC: torch.Tensor | None = None + ecdf_BRiC: torch.Tensor | None = None + if use_chunks: + x_BRiC, nan_ind_BRiC, ecdf_BRiC, y_col_emb_BNE, scaler_stats = ( + self._preprocess_no_group( + rows_RiBC, y, num_train, scaler_cache, task_type + ) + ) + else: + preprocess_fn = ( + self._compiled(self._preprocess_and_group) + if performance_options.enable_torch_compile + else self._preprocess_and_group + ) + x_grouped_BRiCG, y_col_emb_BNE, scaler_stats = preprocess_fn( + rows_RiBC, y, num_train, scaler_cache, task_type + ) + + # --- Phase 1: compute inducing hidden when chunking w/o a pre-built cache. --- + if use_chunks and precomputed_hidden is None: + eff_col_chunk = col_chunk_size if col_chunk_size is not None else C + while True: + try: + precomputed_hidden = self._compute_all_inducing_hidden( + self.feature_distribution_embedder.layers, + x_BRiC[:, :num_train], + nan_ind_BRiC[:, :num_train], + ecdf_BRiC[:, :num_train], + y_col_emb_BNE, + eff_col_chunk, + enable_torch_compile=performance_options.enable_torch_compile, + ) + break + except RuntimeError as e: + if not is_oom_error(e) or eff_col_chunk <= 1: + raise + torch.cuda.empty_cache() + # `torch.mps.empty_cache()` raises where there is no MPS + # backend, which would turn a recoverable OOM into a crash. + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + eff_col_chunk //= 2 + _logger.warning("OOM: halving col_chunk_size to %d", eff_col_chunk) + self.inference_col_chunk_size = eff_col_chunk + + # --- Shared per-chunk loop: embed → dist-embedder → column-aggregator --- + # When not chunking, the single iteration covers all rows. force_recompute_layer + # and return_hidden only apply on the full path (see below). + is_full_path = not use_chunks and precomputed_hidden is None + effective_chunk_size = row_chunk_size if use_chunks else num_rows + + enable_torch_compile = performance_options.enable_torch_compile + process_row_chunk = ( + self._compiled(self._process_row_chunk) + if enable_torch_compile + else self._process_row_chunk + ) + while True: + parts: list[torch.Tensor] = [] + inducing_hidden: list[torch.Tensor] | None = None + try: + for row_chunk_start in range(0, num_rows, effective_chunk_size): + row_chunk_end = min( + row_chunk_start + effective_chunk_size, num_rows + ) + if x_grouped_BRiCG is not None: + x_grouped_chunk = x_grouped_BRiCG[ + :, row_chunk_start:row_chunk_end + ] + else: + x_grouped_chunk = self._group_features( + x_BRiC[:, row_chunk_start:row_chunk_end], + nan_ind_BRiC[:, row_chunk_start:row_chunk_end], + ecdf_BRiC[:, row_chunk_start:row_chunk_end], + ) + if enable_torch_compile: + torch._dynamo.mark_dynamic(x_grouped_chunk, index=0) + torch._dynamo.mark_dynamic(x_grouped_chunk, index=2) + # Will compile two versions: One with dynamic rows and + # one with static rows for the fixed chunk size. + if (row_chunk_end - row_chunk_start) != row_chunk_size: + torch._dynamo.mark_dynamic(x_grouped_chunk, index=1) + + row_embedding_chunk, chunk_hidden = process_row_chunk( + x_grouped_chunk_BRjCG=x_grouped_chunk, + y_col_emb=y_col_emb_BNE, + chunk_start=row_chunk_start, + chunk_end=row_chunk_end, + effective_num_train=effective_num_train, + precomputed_hidden=precomputed_hidden, + save_peak_memory_factor=save_peak_memory_factor, + force_recompute_layer=force_recompute_layer, + return_inducing_hidden=return_inducing_hidden, + is_full_path=is_full_path, + ) + if chunk_hidden is not None: + inducing_hidden = chunk_hidden + parts.append(row_embedding_chunk) + break + except RuntimeError as e: + if not is_oom_error(e) or not use_chunks or effective_chunk_size <= 1: + raise + parts.clear() + torch.cuda.empty_cache() + effective_chunk_size //= 2 + _logger.warning( + "OOM: halving row_chunk_size to %d", effective_chunk_size + ) + self.inference_row_chunk_size = effective_chunk_size + + if use_chunks: + inducing_hidden = precomputed_hidden + x_BRiClE = parts[0] if len(parts) == 1 else torch.cat(parts, dim=1) + return x_BRiClE, inducing_hidden, scaler_stats + + def _process_col_chunk( + self, + *, + x_grouped_chunk_BNCjG: torch.Tensor, + y_col_emb_BNE: torch.Tensor | None, + num_train: int, + ) -> list[torch.Tensor]: + """Compute inducing hidden for one column chunk across all dist-embedder blocks. + + `x_grouped_chunk_BNCjG` has shape `(B, train rows, Cj, G)` — a slice of the + pre-grouped tensor with Cj << C, so the chunked op never sees the full `C` + dim. Returns one `(B, Cj, n_ind, embedding_size)` tensor per block. + """ + B, _, Cj, _ = x_grouped_chunk_BNCjG.shape + + # Embed this column chunk → (B, Rt, Cj, E) + x_emb_BNCjE = self.x_embed(x_grouped_chunk_BNCjG) + E = x_emb_BNCjE.shape[-1] + + # Target-aware y (broadcasts over the Cj columns) + if y_col_emb_BNE is not None and num_train > 0: + x_emb_BNCjE = x_emb_BNCjE + y_col_emb_BNE.unsqueeze(2) + + # (B, Rt, Cj, E) → (B*Cj, Rt, E) + x_flat = x_emb_BNCjE.transpose(1, 2).contiguous().reshape(B * Cj, num_train, E) + + layers = self.feature_distribution_embedder.layers + num_blocks = len(layers) + chunk_outputs: list[torch.Tensor] = [] + for blk_idx, blk in enumerate(layers): + ind = blk.inducing_vectors.unsqueeze(0).expand(B * Cj, -1, -1) + hidden = blk.cross_attn_block1(ind, x_flat) # (B*cc, n_ind, E) + # Reshape for correct batch-column ordering when concatenated + chunk_outputs.append(hidden.reshape(B, Cj, -1, E)) + # Update train embeddings for next block's Step 1 + if blk_idx < num_blocks - 1: + x_flat = blk.cross_attn_block2(x_flat, hidden) + + return chunk_outputs + + def _process_row_chunk( + self, + x_grouped_chunk_BRjCG: torch.Tensor, + y_col_emb: torch.Tensor | None, + chunk_start: int, + chunk_end: int, + effective_num_train: int, + precomputed_hidden: list[torch.Tensor] | None, + save_peak_memory_factor: int | None, + *, + force_recompute_layer: bool, + return_inducing_hidden: bool, + is_full_path: bool, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + """Run one row chunk through dist-embedder and column-aggregator. + + `x_grouped_chunk` has shape `(B, row_chunk_range, C, G)` — a slice + of the pre-grouped tensor. + Returns `(row_embedding_chunk, chunk_hidden)`. `chunk_hidden` is + only non-None when `return_inducing_hidden` is True on the full path. + """ + row_chunk_range = chunk_end - chunk_start + # Number of train rows in this chunk, not overall dataset. + num_train_rows = max(0, min(effective_num_train - chunk_start, row_chunk_range)) + + x_emb = self.x_embed(x_grouped_chunk_BRjCG) + + if y_col_emb is not None and num_train_rows > 0: + y_emb = y_col_emb[:, chunk_start : chunk_start + num_train_rows] + x_emb[:, :num_train_rows] = x_emb[:, :num_train_rows] + y_emb.unsqueeze(2) + + x_emb, chunk_hidden = self.feature_distribution_embedder( + x_BRiCE=x_emb, + num_train_rows=num_train_rows, + cached_hidden=precomputed_hidden, + save_peak_memory_factor=(save_peak_memory_factor if is_full_path else None), + force_recompute_layer=force_recompute_layer and is_full_path, + return_hidden=return_inducing_hidden and is_full_path, + ) + row_embedding_chunk = self.column_aggregator( + x_BRiCE=x_emb, + save_peak_memory_factor=save_peak_memory_factor, + force_recompute_layer=force_recompute_layer and is_full_path, + ) + return row_embedding_chunk, chunk_hidden + + +# --------------------------------------------------------------------------- +# Module interface +# --------------------------------------------------------------------------- + + +def parse_config( + config: dict[str, Any], +) -> tuple[TabPFNV3p5Config, dict[str, Any]]: + """Parse the config dict into a TabPFNV3p5Config, return unused keys.""" + parsed_config = TabPFNV3p5Config(**config) + return parsed_config, parsed_config.get_unused_config(config) + + +def get_architecture( + config: ArchitectureConfig, + *, + cache_trainset_representation: bool = False, +) -> TabPFNV3p5: + """Construct TabPFN v3.5 from the given config.""" + del cache_trainset_representation + assert isinstance(config, TabPFNV3p5Config) + # cache_trainset_representation is accepted for interface compatibility but + # is a no-op: v3.5 uses explicit KV cache passing via forward() parameters + # (kv_cache / return_kv_cache) instead of model-internal caching. + return TabPFNV3p5(config=config) + + +# --------------------------------------------------------------------------- +# Private data utilities +# --------------------------------------------------------------------------- + + +def _prepare_targets( + y: torch.Tensor, + num_train_and_test_rows: int, + batch_size: int, +) -> torch.Tensor: + """Pad y to match num_train_and_test_rows and ensure shape (Ri, B, 1).""" + num_train_labels = y.shape[0] + if num_train_labels > num_train_and_test_rows: + raise ValueError("No test rows provided.") + target_RBT = y.view(num_train_labels, 1 if y.ndim == 1 else batch_size, -1) + return F.pad( + target_RBT, + (0, 0, 0, 0, 0, num_train_and_test_rows - num_train_labels), + value=float("nan"), + ) + + +def _impute_nan_and_inf_with_mean( + x: torch.Tensor, + num_train_rows: int, + scaler_cache: dict[str, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Impute the nan and inf with the mean of the feature. + + Returns: + A tuple of (imputed tensor, is_finite mask). + """ + is_finite = torch.isfinite(x) + if num_train_rows == 0 and scaler_cache is None: + _logging.warning("No training rows or scaler cache provided, imputing with 0.") + if scaler_cache is not None: + feature_means = scaler_cache["mean"] + else: + x_train = torch.where(is_finite[:num_train_rows], x[:num_train_rows], torch.nan) + feature_means = torch.nan_to_num(torch.nanmean(x_train, dim=0), 0) + return torch.where(is_finite, x, feature_means.unsqueeze(0).expand_as(x)), is_finite + + +_ECDF_CONTEXT_KEY = "ecdf_buckets" +"""Key the ECDF ranking context travels under inside the working scaler dict. + +Preprocessing hands one dict back to the caller, so the context rides along with +`mean` and `std`; `TabPFNV3p5Cache` splits it back out into its own field. +""" + +ECDF_CONTEXT_DTYPE: torch.dtype = torch.float32 +"""Storage dtype of the ECDF ranking context.""" + + +_ECDF_CELL_BUDGET = 1 << 23 +"""Cells the ECDF context build and query work on per pass. + +Both carry several intermediates the size of the cells they are given, so a +million-row table done in one pass would cost a multiple of the table itself. +Neither result depends on how the cells are split — the context is built per +column, the ranks per cell — so this only bounds the transients. +""" + + +def _build_ecdf_context( + x_BRiC: torch.Tensor, num_train: int, num_buckets: int +) -> torch.Tensor: + """Summarise the train rows per (batch, column) into ECDF bucket edges. + + Returns `(3, B, C, K)` at `ECDF_CONTEXT_DTYPE`, holding for each of `K = + min(num_buckets, num_train)` edges: the edge value, and the counts of train + values strictly below it and at most equal to it. The counts are exact, so + `_in_context_ecdf` reproduces the true midrank on any value that is an edge. + + A column with at most `K` distinct values gets one edge per distinct value, + which is what makes it exact: consecutive edges then leave no unseen value + between them for interpolation to guess at. Above `K` the edges are spaced + over row positions instead, so every edge carries the same share of the + column and no dense value is skipped. + """ + rows_BNC = x_BRiC[:, :num_train] if num_train > 0 else x_BRiC + columns = rows_BNC.shape[2] + columns_per_pass = max(1, _ECDF_CELL_BUDGET // rows_BNC.shape[1]) + if columns_per_pass < columns: + return torch.cat( + [ + _build_ecdf_context( + x_BRiC[:, :, start : start + columns_per_pass], + num_train, + num_buckets, + ) + for start in range(0, columns, columns_per_pass) + ], + dim=2, + ) + + sorted_BCN = ( + rows_BNC.transpose(1, 2).to(ECDF_CONTEXT_DTYPE).contiguous().sort(dim=-1).values + ) + n = sorted_BCN.shape[-1] + k = min(num_buckets, n) + + # Index of each sorted position within the column's distinct values, so a + # searchsorted over it maps a distinct index back to a row position. + is_new = torch.ones_like(sorted_BCN, dtype=torch.int32) + is_new[..., 1:] = (sorted_BCN[..., 1:] != sorted_BCN[..., :-1]).to(torch.int32) + # In place: on a tall table this is as large as the sorted values themselves. + distinct_idx_BCN = is_new.cumsum_(-1).sub_(1) + num_distinct_BC1 = distinct_idx_BCN[..., -1:] + 1 + + steps = torch.arange(k, device=x_BRiC.device, dtype=ECDF_CONTEXT_DTYPE) + if k > 1: + # Two rulers, because they answer different questions. Evenly spaced + # distinct indices hit every value a column has, but only while it has at + # most `k` of them. Above that the rank error is paid in mass, not in + # distinct values: a column whose rows pile onto a few values inside a + # wide distinct range would starve exactly those values of edges, and + # interpolating across them spans most of the column. Row positions are + # mass-uniform by construction, so they take over there. + distinct_BCK = steps * ((num_distinct_BC1 - 1).to(ECDF_CONTEXT_DTYPE) / (k - 1)) + rows_K = (steps * ((n - 1) / (k - 1))).round().to(torch.int64) + targets_BCK = torch.where( + num_distinct_BC1 <= k, + distinct_BCK.round().to(torch.int32), + distinct_idx_BCN.gather(-1, rows_K.expand(*sorted_BCN.shape[:2], k)), + ).contiguous() + else: + targets_BCK = ( + steps.round().to(torch.int32).expand(*sorted_BCN.shape[:2], k).contiguous() + ) + + below = torch.searchsorted(distinct_idx_BCN, targets_BCK, side="left") + at_most = torch.searchsorted(distinct_idx_BCN, targets_BCK, side="right") + edges_BCK = sorted_BCN.gather(-1, below) + return torch.stack( + [edges_BCK, below.to(ECDF_CONTEXT_DTYPE), at_most.to(ECDF_CONTEXT_DTYPE)] + ) + + +def _ecdf_midrank_counts( + values_BCRi: torch.Tensor, ecdf_context: torch.Tensor +) -> torch.Tensor: + """Midrank of each value against the buckets, as a train-row count.""" + edges_BCK, below_BCK, at_most_BCK = ecdf_context + k = edges_BCK.shape[-1] + # int32 indices halve these two transients; K is a bucket count. + left = torch.searchsorted(edges_BCK, values_BCRi, side="left", out_int32=True) + right = torch.searchsorted(edges_BCK, values_BCRi, side="right", out_int32=True) + + lo_idx = (left - 1).clamp(min=0).to(torch.int64) + hi_idx = left.clamp(max=k - 1).to(torch.int64) + edge_lo = edges_BCK.gather(-1, lo_idx) + edge_hi = edges_BCK.gather(-1, hi_idx) + at_most_lo = at_most_BCK.gather(-1, lo_idx) + below_hi = below_BCK.gather(-1, hi_idx) + + # The two ends coincide only outside the edge range, where the clamp already + # leaves the right count: n above the last edge, and 0 below the first once + # the override below applies. + width = edge_hi - edge_lo + inside = width > 0 + # Divide by 1 outside a bucket instead of masking the quotient: `where` + # backpropagates through the branch it discards, and 0 * inf is NaN. Prompt + # tuning optimises the cells, so that NaN would reach them. + weight = torch.where( + inside, (values_BCRi - edge_lo) / torch.where(inside, width, 1.0), 0.0 + ) + counts = at_most_lo + weight * (below_hi - at_most_lo) + is_edge = right > left + exact = 0.5 * (below_hi + at_most_BCK.gather(-1, hi_idx)) + counts = torch.where(is_edge, exact, counts) + return torch.where((left == 0) & ~is_edge, counts.new_zeros(()), counts) + + +def _in_context_ecdf(x_BRiC: torch.Tensor, ecdf_context: torch.Tensor) -> torch.Tensor: + """Midrank ECDF of each cell value against the train rows, via the buckets. + + A value that is itself a bucket edge gets that edge's exact midrank, which + handles ties. A value inside a bucket is interpolated linearly between the + two counts the bucket's ends bracket — an interval that is empty when the + buckets hold every distinct train value, so the estimate is then exact too. + Inputs must be finite: torch sorts NaN last, so a NaN query would come out at + rank 1.0. + """ + num_rows, columns = x_BRiC.shape[1], x_BRiC.shape[2] + at_most_BCK = ecdf_context[2] + # The last edge is the column maximum, so its at-most count is the row count. + n = at_most_BCK[..., -1:] + + def rank(rows_BRjC: torch.Tensor) -> torch.Tensor: + values_BCRj = rows_BRjC.transpose(1, 2).to(ECDF_CONTEXT_DTYPE).contiguous() + counts_BCRj = _ecdf_midrank_counts(values_BCRj, ecdf_context) + return (counts_BCRj / n).transpose(1, 2).to(x_BRiC.dtype) + + if torch.compiler.is_compiling(): + # A Python loop over the row count would make Dynamo specialise on it, + # which the dynamic-shape marking in `forward` forbids. The compiled path + # ranks all rows in one pass; the inference row chunking bounds them. + return rank(x_BRiC) + + rows_per_pass = max(1, _ECDF_CELL_BUDGET // columns) + ecdf_BRiC = torch.empty(x_BRiC.shape, dtype=x_BRiC.dtype, device=x_BRiC.device) + for start in range(0, num_rows, rows_per_pass): + rows = slice(start, start + rows_per_pass) + ecdf_BRiC[:, rows] = rank(x_BRiC[:, rows]) + return ecdf_BRiC + + +def _ecdf_fourier_features(u: torch.Tensor, num_frequencies: int) -> torch.Tensor: + """Low-frequency sin/cos features of ECDF values in [0, 1]: `(...) -> (..., 2K)`. + + Uses half-period phases (pi * k * u, k = 1..num_frequencies) so u=0 and u=1 + stay distinguishable at every frequency parity (no wrap-around at k=1). + """ + k = torch.arange(1, num_frequencies + 1, device=u.device, dtype=u.dtype) + phase = u.unsqueeze(-1) * (math.pi * k) + return torch.cat([phase.sin(), phase.cos()], dim=-1) + + +def _impute_target_nan_and_inf( + y_NB1: torch.Tensor, + task_type: TaskType, + num_train_rows: int, +) -> torch.Tensor: + # The class imputation for is performed for backwards compatibility. + # We impute the mean and then do a ceil() operation. + # Only apply ceil() to imputed positions to preserve differentiability for + # original values (e.g. during prompt tuning). + y_NB1, is_finite = _impute_nan_and_inf_with_mean(y_NB1, num_train_rows) + if task_type == "regression": + return y_NB1 + return torch.where(is_finite, y_NB1, y_NB1.ceil()) + + +_NAN_INDICATOR = -2.0 +_INFINITY_INDICATOR = 2.0 +_NEG_INFINITY_INDICATOR = 4.0 + + +def _generate_nan_and_inf_indicator(x: torch.Tensor) -> torch.Tensor: + """Generate NaN/Inf indicator features (matches TabPFN v2.5).""" + return ( + torch.isnan(x) * _NAN_INDICATOR + + torch.isposinf(x) * _INFINITY_INDICATOR + + torch.isneginf(x) * _NEG_INFINITY_INDICATOR + ).to(x.dtype) + + +def _safe_log_seqlen( + n: int | torch.Tensor, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """Compute log(n) safely, avoiding fp16 overflow for large `n`.""" + if isinstance(n, torch.Tensor): + return n.to(torch.float32).clamp(min=1).log().to(dtype) + # Materialise `n` via arithmetic on a 0-d tensor rather than + # `torch.as_tensor(n, ...)`. The latter bakes `n` into the graph as a constant and + # emits a `n == ` guard, triggering a recompile on every new value + # `one * n` keeps the value symbolic when `n` is a SymInt. + one = torch.ones((), dtype=torch.float32, device=device) + return (one * n).clamp(min=1).log().to(dtype) + + +def _spline_based_regression_borders(num_buckets: int) -> torch.Tensor: + """Generate hardcoded regression bin borders based on the v2.5 checkpoint. + + Note: Borders are num_buckets + 1! + + + Returns: + An array of shape (num_buckets + 1,) containing the bucket borders. + """ + border_reference_points = [ + (0, -128), + (5, -16.9), + (20, -13), + (100, -9.9), + (200, -8.47), + (500, -6.48), + (1000, -4.40), + ] + # The original model had 5000 buckets. + border_reference_points = ( + border_reference_points + + [(2500, 0)] + + [(5000 - x, -y) for x, y in border_reference_points[::-1]] + ) + x_scale = num_buckets / 5000 + xp = np.array([x for x, _ in border_reference_points]) * x_scale + yp = np.array([y for _, y in border_reference_points]) + return torch.tensor( + np.interp(x=np.arange(num_buckets + 1), xp=xp, fp=yp), dtype=torch.float32 + ) diff --git a/src/tabpfn/browser_auth.py b/src/tabpfn/browser_auth.py index 8588ed32f..e2e1d0c74 100644 --- a/src/tabpfn/browser_auth.py +++ b/src/tabpfn/browser_auth.py @@ -163,6 +163,12 @@ def _get_license_name(hf_repo_id: str) -> str: try: with urllib.request.urlopen(req, timeout=10) as resp: # noqa: S310 data = json.loads(resp.read()) + except urllib.error.HTTPError as exc: + if hf_repo_id == "tabpfn_3_5" and exc.code in (401, 404): + raise TabPFNError( + "TabPFN-3.5 is not publicly available yet. Stay tuned!" + ) from exc + raise TabPFNHuggingFaceGatedRepoError(f"Prior-Labs/{hf_repo_id}") from exc except Exception as exc: raise TabPFNHuggingFaceGatedRepoError(f"Prior-Labs/{hf_repo_id}") from exc license_name = data.get("cardData", {}).get("license_name") diff --git a/src/tabpfn/classifier.py b/src/tabpfn/classifier.py index c0eabd6b8..aae4404fc 100644 --- a/src/tabpfn/classifier.py +++ b/src/tabpfn/classifier.py @@ -628,6 +628,20 @@ def create_default_for_version(cls, version: ModelVersion, **overrides) -> Self: ), "n_estimators": "auto", } + elif version == ModelVersion.V3_5: + options = { + "model_path": prepend_cache_path( + ModelSource.get_v3_5().default_filename + ), + "n_estimators": "auto", + } + elif version == ModelVersion.V3_5_FAST: + options = { + "model_path": prepend_cache_path( + ModelSource.get_v3_5_fast().default_filename + ), + "n_estimators": "auto", + } else: raise ValueError(f"Unknown version: {version}") diff --git a/src/tabpfn/constants.py b/src/tabpfn/constants.py index 9476fb40f..515f03ad2 100644 --- a/src/tabpfn/constants.py +++ b/src/tabpfn/constants.py @@ -35,6 +35,8 @@ class ModelVersion(str, Enum): V2_5 = "v2.5" V2_6 = "v2.6" V3 = "v3" + V3_5 = "v3.5" + V3_5_FAST = "v3.5-fast" NA_PLACEHOLDER = "__MISSING__" diff --git a/src/tabpfn/inference_config.py b/src/tabpfn/inference_config.py index 3b5b31ba9..104e5c56f 100644 --- a/src/tabpfn/inference_config.py +++ b/src/tabpfn/inference_config.py @@ -490,7 +490,8 @@ def raise_if_checkpoints_disagree_on_overridable_fields( def cpu_sample_limit(model_version: ModelVersion) -> int: """Max sample count allowed for CPU inference by default, per model version.""" - return 5000 if model_version == ModelVersion.V3 else 1000 + pre_v3 = (ModelVersion.V2, ModelVersion.V2_5, ModelVersion.V2_6) + return 1000 if model_version in pre_v3 else 5000 def _get_v2_config(preprocessor_configs: list[PreprocessorConfig]) -> InferenceConfig: diff --git a/src/tabpfn/model_loading.py b/src/tabpfn/model_loading.py index afbcb383a..e7c4e160d 100644 --- a/src/tabpfn/model_loading.py +++ b/src/tabpfn/model_loading.py @@ -43,6 +43,8 @@ from tabpfn.settings import settings if TYPE_CHECKING: + from collections.abc import Callable + from sklearn.base import BaseEstimator from tabpfn import TabPFNClassifier, TabPFNRegressor @@ -61,6 +63,8 @@ V_2_5_IDENTIFIER = "v2.5" V_2_6_IDENTIFIER = "v2.6" V_3_IDENTIFIER = "v3" +V_3_5_IDENTIFIER = "v3.5" +V_3_5_FAST_IDENTIFIER = "v3.5-fast" class ModelType(str, Enum): # noqa: D101 @@ -199,32 +203,62 @@ def get_regressor_v3(cls) -> ModelSource: # noqa: D102 filenames=filenames, ) + # From v3.5 on, one checkpoint carries both a classification and a regression + # head, so there is one source per version rather than one per estimator type. -def _get_model_source(version: ModelVersion, model_type: ModelType) -> ModelSource: # noqa: PLR0911 - if version == ModelVersion.V2: - if model_type == ModelType.CLASSIFIER: - return ModelSource.get_classifier_v2() - if model_type == ModelType.REGRESSOR: - return ModelSource.get_regressor_v2() - elif version == ModelVersion.V2_5: - if model_type == ModelType.CLASSIFIER: - return ModelSource.get_classifier_v2_5() - if model_type == ModelType.REGRESSOR: - return ModelSource.get_regressor_v2_5() - elif version == ModelVersion.V2_6: - if model_type == ModelType.CLASSIFIER: - return ModelSource.get_classifier_v2_6() - if model_type == ModelType.REGRESSOR: - return ModelSource.get_regressor_v2_6() - elif version == ModelVersion.V3: - if model_type == ModelType.CLASSIFIER: - return ModelSource.get_classifier_v3() - if model_type == ModelType.REGRESSOR: - return ModelSource.get_regressor_v3() - - raise ValueError( - f"Unsupported version/model combination: {version.value}/{model_type.value}", - ) + @classmethod + def get_v3_5(cls) -> ModelSource: # noqa: D102 + filenames = [ + "tabpfn-v3.5-20260909.safetensors", + "tabpfn-v3.5-20260909_multiclass.safetensors", + ] + return cls( + repo_id="Prior-Labs/tabpfn_3_5", + default_filename="tabpfn-v3.5-20260909.safetensors", + filenames=filenames, + ) + + @classmethod + def get_v3_5_fast(cls) -> ModelSource: # noqa: D102 + # A separate, faster model, not a re-export of `get_v3_5`. + filenames = [ + "tabpfn-v3.5-fast-20260909.safetensors", + ] + return cls( + repo_id="Prior-Labs/tabpfn_3_5", + default_filename="tabpfn-v3.5-fast-20260909.safetensors", + filenames=filenames, + ) + + +def _get_model_source(version: ModelVersion, model_type: ModelType) -> ModelSource: + sources_by_type: dict[ModelType, Callable[[], ModelSource]] | None = { + ModelVersion.V2: { + ModelType.CLASSIFIER: ModelSource.get_classifier_v2, + ModelType.REGRESSOR: ModelSource.get_regressor_v2, + }, + ModelVersion.V2_5: { + ModelType.CLASSIFIER: ModelSource.get_classifier_v2_5, + ModelType.REGRESSOR: ModelSource.get_regressor_v2_5, + }, + ModelVersion.V2_6: { + ModelType.CLASSIFIER: ModelSource.get_classifier_v2_6, + ModelType.REGRESSOR: ModelSource.get_regressor_v2_6, + }, + ModelVersion.V3: { + ModelType.CLASSIFIER: ModelSource.get_classifier_v3, + ModelType.REGRESSOR: ModelSource.get_regressor_v3, + }, + # From v3.5 on, one multitask checkpoint backs both estimator types. + ModelVersion.V3_5: dict.fromkeys(ModelType, ModelSource.get_v3_5), + ModelVersion.V3_5_FAST: dict.fromkeys(ModelType, ModelSource.get_v3_5_fast), + }.get(version) + if sources_by_type is None or model_type not in sources_by_type: + raise ValueError( + "Unsupported version/model combination: " + f"{version.value}/{model_type.value}", + ) + return sources_by_type[model_type]() def _try_huggingface_downloads( @@ -388,6 +422,9 @@ def download_all_models(to: Path) -> None: (ModelVersion.V2_6, ModelSource.get_regressor_v2_6(), "regressor"), (ModelVersion.V3, ModelSource.get_classifier_v3(), "classifier"), (ModelVersion.V3, ModelSource.get_regressor_v3(), "regressor"), + # One multitask checkpoint per v3.5 version backs both estimator types. + (ModelVersion.V3_5, ModelSource.get_v3_5(), "classifier"), + (ModelVersion.V3_5_FAST, ModelSource.get_v3_5_fast(), "classifier"), ]: for ckpt_name in model_source.filenames: path = to / ckpt_name @@ -519,6 +556,8 @@ def _download_model( ModelVersion.V2_5: "tabpfn_2_5", ModelVersion.V2_6: "tabpfn_2_6", ModelVersion.V3: "tabpfn_3", + ModelVersion.V3_5: "tabpfn_3_5", + ModelVersion.V3_5_FAST: "tabpfn_3_5", } if version in _HF_REPOS: try: @@ -591,7 +630,7 @@ def load_model_criterion_config( *, check_bar_distribution_criterion: Literal[False], cache_trainset_representation: bool, - version: Literal["v2", "v2.5", "v2.6", "v3"], + version: Literal["v2", "v2.5", "v2.6", "v3", "v3.5", "v3.5-fast"], estimator_type: Literal["classifier"], download_if_not_exists: bool, softmax_temperature_override: float | None = None, @@ -610,7 +649,7 @@ def load_model_criterion_config( *, check_bar_distribution_criterion: Literal[True], cache_trainset_representation: bool, - version: Literal["v2", "v2.5", "v2.6", "v3"], + version: Literal["v2", "v2.5", "v2.6", "v3", "v3.5", "v3.5-fast"], estimator_type: Literal["regressor"], download_if_not_exists: bool, softmax_temperature_override: float | None = None, @@ -629,7 +668,7 @@ def load_model_criterion_config( check_bar_distribution_criterion: bool, cache_trainset_representation: bool, estimator_type: Literal["regressor", "classifier"], - version: Literal["v2", "v2.5", "v2.6", "v3"], + version: Literal["v2", "v2.5", "v2.6", "v3", "v3.5", "v3.5-fast"], download_if_not_exists: bool, softmax_temperature_override: float | None = None, n_estimators_override: int | None = None, @@ -769,12 +808,17 @@ def _resolve_model_version(model_path: ModelPath | None) -> ModelVersion: if model_path is None: return settings.tabpfn.model_version name = Path(model_path).name - if V_2_6_IDENTIFIER in name: - return ModelVersion.V2_6 - if V_2_5_IDENTIFIER in name: - return ModelVersion.V2_5 - if V_3_IDENTIFIER in name: - return ModelVersion.V3 + # Most specific first: "v3.5-fast" contains "v3.5", which contains "v3". + identifiers = [ + (V_3_5_FAST_IDENTIFIER, ModelVersion.V3_5_FAST), + (V_3_5_IDENTIFIER, ModelVersion.V3_5), + (V_2_6_IDENTIFIER, ModelVersion.V2_6), + (V_2_5_IDENTIFIER, ModelVersion.V2_5), + (V_3_IDENTIFIER, ModelVersion.V3), + ] + for identifier, version in identifiers: + if identifier in name: + return version return ModelVersion.V2 @@ -795,7 +839,7 @@ def resolve_model_version( def resolve_model_path( model_path: ModelPath | list[ModelPath] | None, which: Literal["regressor", "classifier"], - version: Literal["v2", "v2.5", "v2.6", "v3"] = "v3", + version: Literal["v2", "v2.5", "v2.6", "v3", "v3.5", "v3.5-fast"] = "v3", ) -> tuple[ list[Path], list[Path], @@ -814,7 +858,7 @@ def resolve_model_path( interpreted relative to the current working directory. If no file exists there, it falls back to the TabPFN cache directory. which: The type of model ('regressor' or 'classifier'). - version: The model version (currently only 'v2'). + version: The model version, used to pick the default model. Returns: A tuple containing lists of resolved model Path(s), @@ -1323,7 +1367,10 @@ def _resolve_architecture_name(config: ArchitectureConfig) -> str: from tabpfn.architectures.tabpfn_v2_5 import TabPFNV2p5Config # noqa: PLC0415 from tabpfn.architectures.tabpfn_v2_6 import TabPFNV2p6Config # noqa: PLC0415 from tabpfn.architectures.tabpfn_v3 import TabPFNV3Config # noqa: PLC0415 + from tabpfn.architectures.tabpfn_v3_5 import TabPFNV3p5Config # noqa: PLC0415 + if isinstance(config, TabPFNV3p5Config): + return "tabpfn_v3_5" if isinstance(config, TabPFNV3Config): return "tabpfn_v3" if isinstance(config, TabPFNV2p6Config): diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index 5afad1749..061b47fa3 100644 --- a/src/tabpfn/regressor.py +++ b/src/tabpfn/regressor.py @@ -654,6 +654,20 @@ def create_default_for_version(cls, version: ModelVersion, **overrides) -> Self: ), "n_estimators": "auto", } + elif version == ModelVersion.V3_5: + options = { + "model_path": prepend_cache_path( + ModelSource.get_v3_5().default_filename + ), + "n_estimators": "auto", + } + elif version == ModelVersion.V3_5_FAST: + options = { + "model_path": prepend_cache_path( + ModelSource.get_v3_5_fast().default_filename + ), + "n_estimators": "auto", + } else: raise ValueError(f"Unknown version: {version}") diff --git a/tests/test_architectures/test_attention_dispatch.py b/tests/test_architectures/test_attention_dispatch.py index 754d766a9..93b93d36d 100644 --- a/tests/test_architectures/test_attention_dispatch.py +++ b/tests/test_architectures/test_attention_dispatch.py @@ -1,5 +1,5 @@ # Copyright (c) Prior Labs GmbH 2026. -"""TabPFN v3 routes its attention calls through the registry. +"""TabPFN v3 and v3.5 route their attention calls through the registry. Checked from the outside: a registered backend sees the calls it should, described as they really are. @@ -7,12 +7,14 @@ from __future__ import annotations +import inspect from collections.abc import Iterator import pytest import torch -from tabpfn.architectures import tabpfn_v3 +from tabpfn.architectures import tabpfn_v3, tabpfn_v3_5 +from tabpfn.architectures.interface import Architecture, ArchitectureModule from tabpfn.architectures.kv_cache import FP8_KV_DTYPE from tabpfn.architectures.shared import attention_backends from tabpfn.architectures.shared.attention_backends import AttentionSpec @@ -59,21 +61,38 @@ def registry_sandbox() -> Iterator[None]: attention_backends._consult_order = saved_order -def _model(nlayers: int = 2) -> tabpfn_v3.TabPFNV3: - config = tabpfn_v3.TabPFNV3Config( - max_num_classes=10, - num_buckets=5, - embed_dim=48, - nlayers=nlayers, - icl_num_heads=3, - dist_embed_num_heads=3, - feat_agg_num_heads=3, +# Every architecture that dispatches through the registry. +_ARCHITECTURES = [tabpfn_v3, tabpfn_v3_5] + + +def _model(architecture: ArchitectureModule, nlayers: int = 2) -> Architecture: + config, _ = architecture.parse_config( + { + "max_num_classes": 10, + "num_buckets": 5, + "embed_dim": 48, + "nlayers": nlayers, + "icl_num_heads": 3, + "dist_embed_num_heads": 3, + "feat_agg_num_heads": 3, + # Both architectures then run their ICL attention at head_dim 64, in + # one call per layer over train and test rows together. + "feat_agg_num_cls_tokens": 4, + "icl_num_kv_heads_test": None, + } ) - model = tabpfn_v3.get_architecture(config, cache_trainset_representation=False) + model = architecture.get_architecture(config, cache_trainset_representation=False) model.to(torch.float32) return model +def _forward_kwargs(model: Architecture) -> dict[str, str]: + """v3.5 takes the task per call; v3 is built for one task.""" + if "task_type" in inspect.signature(model.forward).parameters: + return {"task_type": "multiclass"} + return {} + + def _inputs() -> tuple[torch.Tensor, torch.Tensor]: torch.manual_seed(0) x = torch.randn(20, 2, 5, dtype=torch.float32) * 0.1 @@ -82,15 +101,18 @@ def _inputs() -> tuple[torch.Tensor, torch.Tensor]: @pytest.mark.usefixtures("registry_sandbox") +@pytest.mark.parametrize("architecture", _ARCHITECTURES, ids=lambda m: m.__name__) @torch.no_grad() -def test_backend_receives_the_calls_it_prefers() -> None: +def test_backend_receives_the_calls_it_prefers( + architecture: ArchitectureModule, +) -> None: """A backend preferring the ICL shape runs once per ICL layer.""" backend = _RecordingBackend(take=_is_icl_spec) attention_backends.register_attention_backend(backend) - model = _model(nlayers=2) + model = _model(architecture, nlayers=2) x, y = _inputs() - model(x, y) + model(x, y, **_forward_kwargs(model)) assert backend.runs == 2 # one per ICL layer icl_specs = [spec for spec in backend.specs if _is_icl_spec(spec)] @@ -103,19 +125,22 @@ def test_backend_receives_the_calls_it_prefers() -> None: @pytest.mark.usefixtures("registry_sandbox") +@pytest.mark.parametrize("architecture", _ARCHITECTURES, ids=lambda m: m.__name__) @torch.no_grad() -def test_cached_predict_specs_report_the_quantized_cache() -> None: +def test_cached_predict_specs_report_the_quantized_cache( + architecture: ArchitectureModule, +) -> None: """On the cache path the specs carry the stored KV dtype.""" backend = _RecordingBackend() # observe only attention_backends.register_attention_backend(backend) - model = _model() + model = _model(architecture) x, y = _inputs() - _, cache = model(x, y, return_kv_cache=True) + _, cache = model(x, y, return_kv_cache=True, **_forward_kwargs(model)) cache = cache.quantize(FP8_KV_DTYPE) backend.specs.clear() - model(x[10:], y, kv_cache=cache, x_is_test_only=True) + model(x[10:], y, kv_cache=cache, x_is_test_only=True, **_forward_kwargs(model)) quantized = [s for s in backend.specs if s.quantized_kv_dtype is not None] assert quantized, "no spec reported the quantized cache" diff --git a/tests/test_architectures/test_compile.py b/tests/test_architectures/test_compile.py index e73199fe0..1602b6982 100644 --- a/tests/test_architectures/test_compile.py +++ b/tests/test_architectures/test_compile.py @@ -14,6 +14,7 @@ from __future__ import annotations +import inspect import logging import re from pathlib import Path @@ -23,8 +24,12 @@ from torch.torch_version import TorchVersion import tabpfn -from tabpfn.architectures import tabpfn_v3 -from tabpfn.architectures.interface import PerformanceOptions +from tabpfn.architectures import tabpfn_v3, tabpfn_v3_5 +from tabpfn.architectures.interface import ( + Architecture, + ArchitectureModule, + PerformanceOptions, +) # This test reads Dynamo's graph-break log, whose wording has changed over # torch releases (verified on 2.12 and 2.13). Older builds phrase it @@ -58,23 +63,40 @@ _TABPFN_ROOT = str(Path(tabpfn.__file__).parent) -def _tiny_model() -> tabpfn_v3.TabPFNV3: - config = tabpfn_v3.TabPFNV3Config( - max_num_classes=10, - num_buckets=5, - embed_dim=48, - nlayers=1, - icl_num_heads=3, - dist_embed_num_heads=3, - feat_agg_num_heads=3, +# Every architecture that honours `enable_torch_compile`. +_ARCHITECTURES = [tabpfn_v3, tabpfn_v3_5] + + +def _tiny_model(architecture: ArchitectureModule) -> Architecture: + config, _ = architecture.parse_config( + { + "max_num_classes": 10, + "num_buckets": 5, + "embed_dim": 48, + "nlayers": 1, + "icl_num_heads": 3, + "dist_embed_num_heads": 3, + "feat_agg_num_heads": 3, + # Both architectures then run their ICL attention at head_dim 64, in + # one call per layer over train and test rows together. + "feat_agg_num_cls_tokens": 4, + "icl_num_kv_heads_test": None, + } ) - model = tabpfn_v3.get_architecture(config, cache_trainset_representation=False) + model = architecture.get_architecture(config, cache_trainset_representation=False) model.to(torch.float32) return model +def _forward_kwargs(model: Architecture) -> dict[str, str]: + """v3.5 takes the task per call; v3 is built for one task.""" + if "task_type" in inspect.signature(model.forward).parameters: + return {"task_type": "multiclass"} + return {} + + def _compiled_forward_graph_breaks( - model: tabpfn_v3.TabPFNV3, + model: Architecture, x: torch.Tensor, y: torch.Tensor, monkeypatch: pytest.MonkeyPatch, @@ -104,7 +126,10 @@ def emit(self, record: logging.LogRecord) -> None: torch._dynamo.utils.counters.clear() with torch.no_grad(): out = model( - x, y, performance_options=PerformanceOptions(enable_torch_compile=True) + x, + y, + performance_options=PerformanceOptions(enable_torch_compile=True), + **_forward_kwargs(model), ) finally: logging.getLogger("torch._dynamo").removeHandler(handler) @@ -114,7 +139,9 @@ def emit(self, record: logging.LogRecord) -> None: @torch.no_grad() +@pytest.mark.parametrize("architecture", _ARCHITECTURES, ids=lambda m: m.__name__) def test__enable_torch_compile__no_graph_break_in_tabpfn_code( + architecture: ArchitectureModule, monkeypatch: pytest.MonkeyPatch, ) -> None: """Compiled regions must not break on tabpfn's own code. @@ -124,12 +151,12 @@ def test__enable_torch_compile__no_graph_break_in_tabpfn_code( compiled stage. The output stays correct, so only the break log shows it. """ - model = _tiny_model() + model = _tiny_model(architecture) torch.manual_seed(0) x = torch.randn(30, 2, 5, dtype=torch.float32) * 0.1 y = torch.randint(0, 10, [15, 2], dtype=torch.float32) - out_eager = model(x, y) + out_eager = model(x, y, **_forward_kwargs(model)) out_compiled, records = _compiled_forward_graph_breaks(model, x, y, monkeypatch) torch.testing.assert_close(out_compiled, out_eager, atol=1e-5, rtol=1e-5) diff --git a/tests/test_architectures/test_kv_cache.py b/tests/test_architectures/test_kv_cache.py index 048e166b8..0a0a8de97 100644 --- a/tests/test_architectures/test_kv_cache.py +++ b/tests/test_architectures/test_kv_cache.py @@ -13,13 +13,13 @@ import torch from tabpfn import TabPFNClassifier +from tabpfn.architectures import tabpfn_v3, tabpfn_v3_5 from tabpfn.architectures.kv_cache import ( FP8_KV_DTYPE, KVCacheEntry, _dequantize_tensor, _quantize_tensor, ) -from tabpfn.architectures.tabpfn_v3 import TabPFNV3Config, get_cache_size from tabpfn.inference import _resolve_kv_cache_precision @@ -130,13 +130,25 @@ def test_resolve_rejects_fp8_on_mps() -> None: ) -def test_get_cache_size_accepts_fp8() -> None: - config = TabPFNV3Config() +@pytest.mark.parametrize( + ("get_cache_size", "config", "extra_kwargs"), + [ + (tabpfn_v3.get_cache_size, tabpfn_v3.TabPFNV3Config(), {}), + # v3.5 is multitask, so the cache contents depend on the task. + ( + tabpfn_v3_5.get_cache_size, + tabpfn_v3_5.TabPFNV3p5Config(), + {"task_type": "multiclass"}, + ), + ], +) +def test_get_cache_size_accepts_fp8(get_cache_size, config, extra_kwargs) -> None: kwargs = { "n_train": 1000, "n_features": 20, "model_config": config, "base_dtype": torch.float32, + **extra_kwargs, } # int8 and fp8 are both one byte per element plus scales. assert get_cache_size(kv_cache_precision="fp8", **kwargs) == get_cache_size( diff --git a/tests/test_architectures/test_tabpfn_v3_5.py b/tests/test_architectures/test_tabpfn_v3_5.py new file mode 100644 index 000000000..23f94ed8c --- /dev/null +++ b/tests/test_architectures/test_tabpfn_v3_5.py @@ -0,0 +1,1177 @@ +# Copyright (c) Prior Labs GmbH 2026. + +"""Tests for the v3.5 single-file model. + +v3.5 is a multitask model: `task_type` is a per-`forward()` argument, so one +instance, and one checkpoint, handles both multiclass and regression. It ranks +cells against `cell_ecdf_num_buckets` bucket edges per column rather than +against every train row, so the inference cache stops growing with the table. +""" + +from __future__ import annotations + +import ast +import dataclasses +import functools +import inspect +import sys +from typing import Literal + +import numpy as np +import pytest +import torch + +from tabpfn import TabPFNClassifier +from tabpfn.architectures import tabpfn_v3_5 +from tabpfn.architectures.interface import PerformanceOptions +from tabpfn.architectures.kv_cache import ( + FP8_KV_DTYPE, + QUANTIZED_KV_DTYPE, + KVCacheEntry, + QuantizedKVCacheEntry, +) +from tabpfn.architectures.tabpfn_v3_5 import ( + TabPFNV3p5, + TabPFNV3p5Cache, + TabPFNV3p5Config, + get_cache_size, +) +from tabpfn.constants import ModelVersion, TaskType +from tabpfn.utils import get_autocast_context + +MAX_NUM_CLASSES = 5 +NUM_TRAIN, NUM_TEST, BATCH, NUM_FEATURES = 20, 4, 2, 5 +TASK_TYPES: list[TaskType] = ["multiclass", "regression"] + +# Shrunk to keep the tests fast; every stage of the model is still exercised. +_SMALL_CONFIG: dict[str, object] = { + "max_num_classes": MAX_NUM_CLASSES, + "num_buckets": 32, + "embed_dim": 32, + "nlayers": 2, + "icl_num_heads": 4, + "icl_num_kv_heads_test": 1, + "dist_embed_num_heads": 4, + "dist_embed_num_blocks": 1, + "feat_agg_num_heads": 4, + "feat_agg_num_blocks": 1, + "feat_agg_num_cls_tokens": 2, + "dist_embed_num_inducing_points": 8, + # Small enough that the chunked inference path splits the test inputs. + "inference_row_chunk_size": 8, + "inference_col_chunk_size": 2, +} + + +def _config(config_overrides: dict[str, object]) -> TabPFNV3p5Config: + config, _unused = tabpfn_v3_5.parse_config({**_SMALL_CONFIG, **config_overrides}) + return config + + +def _get_model(**config_overrides: object) -> TabPFNV3p5: + """A small v3.5 model in eval mode, with no all-zero parameters.""" + config = _config(config_overrides) + arch = tabpfn_v3_5.get_architecture(config, cache_trainset_representation=False) + # Several modules zero-init their residual out-projections; a fully-zero + # projection masks its sublayer and would hide a bug in it. + gen = torch.Generator().manual_seed(0) + with torch.no_grad(): + for param in arch.parameters(): + if param.numel() > 0 and bool((param == 0).all()): + param.normal_(std=0.02, generator=gen) + arch.to(torch.float32) + return arch.eval() + + +def _inputs( + task_type: TaskType, + *, + n_train_classes: int = MAX_NUM_CLASSES, + batch: int = BATCH, +) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(0) + x = torch.randn(NUM_TRAIN + NUM_TEST, batch, NUM_FEATURES) * 0.1 + if task_type == "regression": + return x, torch.randn(NUM_TRAIN, batch) + y = torch.arange(NUM_TRAIN).unsqueeze(1).repeat(1, batch) % n_train_classes + return x, y.float() + + +def _assert_outputs_equal( + actual: dict[str, torch.Tensor], + expected: dict[str, torch.Tensor], + *, + atol: float, +) -> None: + assert actual.keys() == expected.keys(), "Output keys do not match" + for key, value in expected.items(): + assert torch.allclose(actual[key], value, atol=atol), ( + f"Outputs for '{key}' do not match." + ) + + +# --------------------------------------------------------------------------- +# Config and module +# --------------------------------------------------------------------------- + + +def test__config__defaults__match_the_v3_5_checkpoint() -> None: + """The defaults rebuild the released checkpoint's architecture. + + The two head sizes stay unset: they come from the checkpoint, and the base + `ArchitectureConfig` leaves them at -1. + """ + config = TabPFNV3p5Config() + assert dataclasses.asdict(config) == { + "name": "TabPFN-v3.5", + "max_num_classes": -1, + "num_buckets": -1, + "embed_dim": 128, + "dist_embed_num_blocks": 3, + "dist_embed_num_heads": 8, + "dist_embed_num_inducing_points": 128, + "feature_group_size": 3, + "feat_agg_num_blocks": 3, + "feat_agg_num_heads": 8, + "feat_agg_num_cls_tokens": 8, + "feat_agg_rope_base": 100_000, + "nlayers": 24, + "icl_num_heads": 16, + "icl_num_kv_heads": None, + "icl_num_kv_heads_test": 1, + "decoder_head_dim": 64, + "decoder_num_heads": 6, + "decoder_use_softmax_scaling": True, + "ff_factor": 2, + "softmax_scaling_mlp_hidden_dim": 64, + "fourier_encoding_num_frequencies": 32, + "cell_ecdf_num_frequencies": 4, + "cell_ecdf_num_buckets": 8192, + "cell_embed_row_chunk_size": 2048, + "inference_row_chunk_size": 2048, + "inference_col_chunk_size": 4, + } + + +def test__module_imports__only_tabpfn_and_third_party() -> None: + """The architecture must stay self-contained within the tabpfn package.""" + tree = ast.parse(inspect.getsource(tabpfn_v3_5)) + roots = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + roots.update(alias.name.split(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + roots.add(node.module.split(".")[0]) + assert roots == { + "__future__", + "collections", + "contextlib", + "dataclasses", + "functools", + "logging", + "math", + "numpy", + "pydantic", + "tabpfn", + "torch", + "typing", + "typing_extensions", + } + + +def test__parse_config__training_only_keys__reported_as_unused() -> None: + """Training checkpoints carry loss and muP keys; v3.5 must ignore them.""" + training_only = {"enable_mup": True, "weight_regression_ce_loss": 1.0} + _config, unused = tabpfn_v3_5.parse_config({**_SMALL_CONFIG, **training_only}) + assert set(unused) == set(training_only) + + +def test__get_supported_kv_cache_precisions__advertises_the_quantized_dtypes() -> None: + """The engine resolves to "auto" unless the architecture lists its dtypes.""" + assert _get_model().get_supported_kv_cache_precisions() == ("auto", "int8", "fp8") + + +# --------------------------------------------------------------------------- +# Forward pass +# --------------------------------------------------------------------------- + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__forward__output_shapes(task_type: TaskType) -> None: + arch = _get_model() + x, y = _inputs(task_type) + out = arch(x, y, task_type=task_type) + width = ( + MAX_NUM_CLASSES if task_type == "multiclass" else _SMALL_CONFIG["num_buckets"] + ) + assert out.shape == (NUM_TEST, BATCH, width) + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__forward__only_return_standard_out_false__returns_embeddings( + task_type: TaskType, +) -> None: + """v3.5 carries no losses, so the dict output holds only the three tensors.""" + arch = _get_model() + x, y = _inputs(task_type) + output = arch(x, y, task_type=task_type, only_return_standard_out=False) + assert set(output) == {"standard", "train_embeddings", "test_embeddings"} + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__forward_pass_equal_with_save_peak_memory_enabled_and_disabled( + task_type: TaskType, +) -> None: + arch = _get_model() + x, y = _inputs(task_type) + + without = arch(x, y, task_type=task_type, only_return_standard_out=False) + with_saving = arch( + x, + y, + task_type=task_type, + only_return_standard_out=False, + performance_options=PerformanceOptions(save_peak_memory_factor=4), + ) + _assert_outputs_equal(with_saving, without, atol=1e-6) + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__forward_pass_equal_with_checkpointing_enabled_and_disabled( + task_type: TaskType, +) -> None: + arch = _get_model() + x, y = _inputs(task_type) + + without = arch(x, y, task_type=task_type, only_return_standard_out=False) + with_recompute = arch( + x, + y, + task_type=task_type, + only_return_standard_out=False, + performance_options=PerformanceOptions(force_recompute_layer=True), + ) + _assert_outputs_equal(with_recompute, without, atol=1e-6) + + +@torch.no_grad() +def test__batch_size_one__nan_and_inf_in_features__still_works() -> None: + arch = _get_model() + x = torch.randn(100, 1, 1, dtype=torch.float32) * 0.1 + x[10, 0] = float("nan") + x[11, 0] = float("inf") + y = torch.randint(0, MAX_NUM_CLASSES, [97, 1], dtype=torch.float32) + + output = arch(x, y, task_type="multiclass") + + assert output.shape == (3, 1, MAX_NUM_CLASSES) + assert torch.isfinite(output).all() + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__forward__no_test_set_works_batch_size_one(task_type: TaskType) -> None: + arch = _get_model() + x = torch.randn(1, 1, NUM_FEATURES, dtype=torch.float32) * 0.1 + y = torch.randint(0, MAX_NUM_CLASSES, [1, 1], dtype=torch.float32) + + out = arch(x, y, task_type=task_type, only_return_standard_out=False) + + assert out["standard"].shape[:2] == (0, 1) + + +@torch.no_grad() +@pytest.mark.parametrize( + "invalid_target", + [-1.0, -0.5, MAX_NUM_CLASSES - 0.5, MAX_NUM_CLASSES, -np.inf, np.inf], +) +def test__forward__multiclass_target_out_of_range__raises( + invalid_target: float, +) -> None: + arch = _get_model() + x, y = _inputs("multiclass") + y[0, 0] = invalid_target + with pytest.raises(ValueError, match="Target is out of range"): + arch(x, y, task_type="multiclass") + + +@torch.no_grad() +def test__forward__nan_and_highest_class_in_second_batch__cached_matches_uncached() -> ( + None +): + arch = _get_model() + x, y = _inputs("multiclass", n_train_classes=2) + y[0, 0] = np.nan + y[0, 1] = MAX_NUM_CLASSES - 1 + expected = arch(x, y, task_type="multiclass") + + _, cache = arch(x[:NUM_TRAIN], y, task_type="multiclass", return_kv_cache=True) + actual = arch( + x[NUM_TRAIN:], + y, + task_type="multiclass", + kv_cache=cache, + x_is_test_only=True, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__chunked_inference_matches_standard_forward(task_type: TaskType) -> None: + """The row/column chunking must not change the prediction.""" + arch = _get_model() + x, y = _inputs(task_type) + options = arch.get_default_performance_options() + assert options.use_chunkwise_inference + + standard = arch( + x, + y, + task_type=task_type, + only_return_standard_out=False, + performance_options=PerformanceOptions(use_chunkwise_inference=False), + ) + chunked = arch( + x, + y, + task_type=task_type, + only_return_standard_out=False, + performance_options=options, + ) + _assert_outputs_equal(chunked, standard, atol=1e-5) + + +@torch.no_grad() +def test__chunked_inference_recovers_from_oom(monkeypatch: pytest.MonkeyPatch) -> None: + """A recoverable OOM during chunked inference must not crash the forward. + + The column-chunk handler reacts to an OOM by freeing memory, halving the + chunk and retrying. The recovered output must match the standard forward. + """ + arch = _get_model() + x, y = _inputs("multiclass") + expected = arch(x, y, task_type="multiclass", only_return_standard_out=False) + + # Raise a single OOM the first time a column chunk is processed, so the handler + # must free memory, halve the column chunk and retry. Patched on the class so + # the bound method still exposes `__func__` for `_compiled`. + original_process_col_chunk = TabPFNV3p5._process_col_chunk + calls = {"n": 0} + + def _process_col_chunk_oom_once( + self: TabPFNV3p5, *args: object, **kwargs: object + ) -> object: + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("CUDA out of memory (simulated)") + return original_process_col_chunk(self, *args, **kwargs) + + monkeypatch.setattr(TabPFNV3p5, "_process_col_chunk", _process_col_chunk_oom_once) + + recovered = arch( + x, + y, + task_type="multiclass", + only_return_standard_out=False, + performance_options=PerformanceOptions(use_chunkwise_inference=True), + ) + + assert calls["n"] > 1, "the simulated OOM never triggered a retry" + _assert_outputs_equal(recovered, expected, atol=1e-5) + + +# --------------------------------------------------------------------------- +# Many-class decoder +# --------------------------------------------------------------------------- + + +@torch.no_grad() +def test__forward_many_class_head__fewer_classes_than_max__pads_absent_columns() -> ( + None +): + """Classes missing from the train targets still get a column, at the zero logit. + + The decoder narrows its one-hot to the classes present and pads the output + back, so the absent columns must carry the logit of a zero attention output. + """ + arch = _get_model() + n_train_classes = 2 + x, y = _inputs("multiclass", n_train_classes=n_train_classes) + + out = arch(x, y, task_type="multiclass") + + assert out.shape[-1] == MAX_NUM_CLASSES + zero_logit = float(np.log(1e-5 + 3e-5)) + assert torch.allclose( + out[..., n_train_classes:], + torch.full_like(out[..., n_train_classes:], zero_logit), + ) + + +@torch.no_grad() +@pytest.mark.skipif(sys.platform == "win32", reason="float64 tests fail on Windows") +def test__many_class_decoder__unused_classes__matches_full_width_one_hot() -> None: + """Narrowing the one-hot to the present classes must not change the output. + + `head_dim` below the class count makes the full-width reference span three + folded attention passes where the narrowed path needs one, so a mismatch in + the trim, the chunking or the padding surfaces here. + """ + num_classes, input_size, num_heads, head_dim = 10, 12, 3, 4 + batch, num_train, num_test = 2, 17, 3 + + torch.manual_seed(42) + decoder = tabpfn_v3_5.ManyClassDecoder( + max_num_classes=num_classes, + input_size=input_size, + head_dim=head_dim, + num_heads=num_heads, + ).to(torch.float64) + train_emb = torch.randn(batch, num_train, input_size, dtype=torch.float64) + test_emb = torch.randn(batch, num_test, input_size, dtype=torch.float64) + # Only classes 0..2 occur, so the decoder trims 10 columns down to 3. + targets = (torch.arange(num_train) % 3).repeat(batch, 1).to(torch.float64) + + train_keys = decoder.project_keys(train_emb) + actual = decoder(train_keys, test_emb, targets, num_present_classes=3) + + q_BMHD = decoder.q_projection(test_emb).view(batch, num_test, num_heads, head_dim) + one_hot_BNHT = ( + torch.nn.functional.one_hot(targets.long(), num_classes=num_classes) + .to(torch.float64) + .unsqueeze(2) + .expand(-1, -1, num_heads, -1) + .contiguous() + ) + reference_BMT = tabpfn_v3_5._chunked_class_attention( + q_BMHD.contiguous(), train_keys, one_hot_BNHT + ).mean(2) + expected = torch.log(torch.clamp(reference_BMT.transpose(0, 1), min=1e-5) + 3e-5) + + assert actual.shape == (num_test, batch, num_classes) + assert torch.allclose(actual, expected, atol=1e-12), ( + f"max abs diff: {(actual - expected).abs().max()}" + ) + + +# --------------------------------------------------------------------------- +# In-context ECDF +# --------------------------------------------------------------------------- + + +def _exact_midranks(x_BRiC: torch.Tensor, num_train: int) -> torch.Tensor: + """Midrank ECDF against every train row, the definition v3.5 approximates.""" + sorted_BCN = x_BRiC[:, :num_train].transpose(1, 2).contiguous().sort(dim=-1).values + values_BCRi = x_BRiC.transpose(1, 2).contiguous() + left = torch.searchsorted(sorted_BCN, values_BCRi, side="left") + right = torch.searchsorted(sorted_BCN, values_BCRi, side="right") + return (0.5 * (left + right).float() / num_train).transpose(1, 2) + + +def _bucketed_midranks( + x_BRiC: torch.Tensor, num_train: int, num_buckets: int +) -> torch.Tensor: + context = tabpfn_v3_5._build_ecdf_context(x_BRiC, num_train, num_buckets) + assert context.dtype == tabpfn_v3_5.ECDF_CONTEXT_DTYPE + batch, _rows, columns = x_BRiC.shape + assert context.shape == (3, batch, columns, min(num_buckets, num_train)) + return tabpfn_v3_5._in_context_ecdf(x_BRiC, context) + + +@pytest.mark.parametrize( + ("kind", "num_buckets", "column"), + [ + # Fewer train rows than buckets: the buckets are the rows. + ("all-rows-fit", 10_000, torch.linspace(-3.0, 3.0, 4000)), + # More rows than buckets, but few enough distinct values to keep them all. + ("low-cardinality", 100, torch.randint(0, 7, (4000,)).float()), + # A value seen once still gets its own bucket edge. + ("one-rare-value", 100, torch.cat([torch.zeros(3999), torch.ones(1)])), + # A constant column: every row lands on the single edge. + ("constant", 100, torch.full((4000,), 4.2)), + ], +) +def test__in_context_ecdf__buckets_cover_every_value__matches_exact_midranks( + kind: str, num_buckets: int, column: torch.Tensor +) -> None: + """Where no distinct value is dropped, bucketing must change nothing at all.""" + del kind + x_BRiC = column.reshape(1, -1, 1) + ranks = _bucketed_midranks(x_BRiC, x_BRiC.shape[1], num_buckets) + assert torch.equal(ranks, _exact_midranks(x_BRiC, x_BRiC.shape[1])) + + +def test__build_ecdf_context__dense_values_in_a_wide_range__keep_their_rank() -> None: + """Edges must follow the rows, not the distinct values, once they run out. + + Half the rows sit on six values holding six of ~10 000 distinct indices, so + spacing the edges over distinct values skips every one of them and ranks half + the column inside one bucket. + """ + torch.manual_seed(0) + num_rows, num_buckets = 20_000, 512 + dense = torch.arange(6.0).repeat_interleave(num_rows // 12) + tail = torch.rand(num_rows - dense.numel()) * 1000 + 10 + # The six dense values ride along as test rows so they get ranked too. + x_BRiC = torch.cat([dense, tail, torch.arange(6.0)]).reshape(1, -1, 1) + + context = tabpfn_v3_5._build_ecdf_context(x_BRiC, num_rows, num_buckets) + ranks = tabpfn_v3_5._in_context_ecdf(x_BRiC, context) + expected = _exact_midranks(x_BRiC, num_rows) + assert (ranks[:, num_rows:] - expected[:, num_rows:]).abs().max() < 1 / num_buckets + + +@pytest.mark.parametrize(("num_buckets", "tolerance"), [(1000, 1e-3), (100, 1e-2)]) +def test__in_context_ecdf__more_distinct_values_than_buckets__error_below_one_bucket( + num_buckets: int, tolerance: float +) -> None: + """Dropped values cost at most a bucket's worth of rank, on any column scale.""" + torch.manual_seed(0) + num_rows = 20_000 + scales = [1.0, 1e3, 1e5, 1e8, 1e-4, 1e30, 1e-30] + x_BRiC = torch.stack([torch.randn(num_rows) * s for s in scales], dim=-1).unsqueeze( + 0 + ) + ranks = _bucketed_midranks(x_BRiC, num_rows, num_buckets) + assert torch.isfinite(ranks).all() + assert (ranks - _exact_midranks(x_BRiC, num_rows)).abs().max() < tolerance + + +def test__build_ecdf_context__column_chunking__builds_the_same_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The context is built per column, so the chunk width must not matter.""" + torch.manual_seed(0) + x_BRiC = torch.randn(2, 200, 7) + one_pass = tabpfn_v3_5._build_ecdf_context(x_BRiC, 150, 32) + # 200 rows over a 300-cell budget gives one column per pass, not a divisor of 7. + monkeypatch.setattr(tabpfn_v3_5, "_ECDF_CELL_BUDGET", 300) + chunked = tabpfn_v3_5._build_ecdf_context(x_BRiC, 150, 32) + assert chunked.shape == one_pass.shape + assert torch.equal(chunked, one_pass) + + +@pytest.mark.parametrize( + ("kind", "num_train", "num_buckets"), + [ + ("every value an edge", 40, 8192), + ("interpolating", 300, 8), + ("duplicate edges", 300, 64), + ], +) +def test__in_context_ecdf__cells_under_optimisation__keep_a_finite_gradient( + kind: str, num_train: int, num_buckets: int +) -> None: + """Prompt tuning optimises the cells, so the ranks must stay differentiable.""" + torch.manual_seed(0) + values = ( + torch.randint(0, 4, (1, 400, 3)).float() + if kind == "duplicate edges" + else torch.randn(1, max(num_train, 400), 3) + ) + values.requires_grad_() + context = tabpfn_v3_5._build_ecdf_context(values, num_train, num_buckets) + tabpfn_v3_5._in_context_ecdf(values, context).sum().backward() + assert torch.isfinite(values.grad).all() + + +def test__in_context_ecdf__row_chunking__does_not_move_a_single_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ranking is per cell, so the chunk boundaries must not show up anywhere.""" + torch.manual_seed(0) + x_BRiC = torch.randn(2, 500, 3) + context = tabpfn_v3_5._build_ecdf_context(x_BRiC, 400, 64) + one_pass = tabpfn_v3_5._in_context_ecdf(x_BRiC, context) + # Small enough to split the 500 rows many times, and not a divisor of them. + monkeypatch.setattr(tabpfn_v3_5, "_ECDF_CELL_BUDGET", 21) + assert torch.equal(tabpfn_v3_5._in_context_ecdf(x_BRiC, context), one_pass) + + +def test__in_context_ecdf__value_inside_a_bucket__is_interpolated_not_snapped() -> None: + """Within a bucket the rank rises with the value instead of stepping.""" + x_BRiC = torch.arange(1000.0).reshape(1, -1, 1) + ranks = _bucketed_midranks(x_BRiC, 1000, 10).flatten() + assert (ranks[1:] > ranks[:-1]).all() + # On a uniform column linear interpolation recovers the exact ranks. + assert (ranks - _exact_midranks(x_BRiC, 1000).flatten()).abs().max() < 1e-3 + + +@torch.no_grad() +@pytest.mark.parametrize("cell_ecdf_num_buckets", [8192, 8]) +def test__preprocess_raw__cached_context__ranks_test_rows_identically( + cell_ecdf_num_buckets: int, +) -> None: + """A cached run must rank test rows against exactly the buckets it stored. + + 8 buckets over 20 train rows is lossy, which is what makes this fail if either + path builds its own context. + """ + arch = _get_model(cell_ecdf_num_buckets=cell_ecdf_num_buckets) + x, _ = _inputs("multiclass") + _, _, ecdf_full, scaler_stats = arch._preprocess_raw(x, num_train=NUM_TRAIN) + _, _, ecdf_from_cache, _ = arch._preprocess_raw( + x[NUM_TRAIN:], num_train=0, scaler_cache=scaler_stats + ) + assert torch.equal(ecdf_full[:, NUM_TRAIN:], ecdf_from_cache) + + +# --------------------------------------------------------------------------- +# KV cache +# --------------------------------------------------------------------------- + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +@pytest.mark.parametrize("use_chunkwise", [False, True]) +# 16 buckets over 20 train rows makes the cached ranks lossy, so the cached and +# uncached paths agree only if both rank against the same buckets. +@pytest.mark.parametrize("cell_ecdf_num_buckets", [10_000, 16]) +def test__kv_cache__matches_standard_forward( + task_type: TaskType, use_chunkwise: bool, cell_ecdf_num_buckets: int +) -> None: + """Reusing the cache on test-only rows must reproduce the full forward. + + Not bit-for-bit: the cached call feeds the attention kernels a test-only + sequence instead of train+test, and the kernel's reduction order follows the + sequence length. The tolerance is a float-noise bound. + """ + arch = _get_model(cell_ecdf_num_buckets=cell_ecdf_num_buckets) + x, y = _inputs(task_type) + perf = PerformanceOptions(use_chunkwise_inference=use_chunkwise) + + out_standard = arch(x, y, task_type=task_type, performance_options=perf) + out_store, cache = arch( + x, y, task_type=task_type, performance_options=perf, return_kv_cache=True + ) + + assert isinstance(cache, TabPFNV3p5Cache) + assert not cache.is_empty() + assert len(cache.kv) == _SMALL_CONFIG["nlayers"] + assert cache.train_shape == (BATCH, NUM_TRAIN) + torch.testing.assert_close(out_store, out_standard, rtol=0, atol=1e-6) + + # Test-only rows against the cache, and the full tensor against the cache. + out_test_only = arch( + x[NUM_TRAIN:], + y, + task_type=task_type, + performance_options=perf, + kv_cache=cache, + x_is_test_only=True, + ) + out_full = arch(x, y, task_type=task_type, performance_options=perf, kv_cache=cache) + torch.testing.assert_close(out_test_only, out_standard, rtol=0, atol=1e-5) + torch.testing.assert_close(out_full, out_standard, rtol=0, atol=1e-5) + + +@torch.no_grad() +def test__kv_cache__x_is_test_only_without_cache__raises() -> None: + arch = _get_model() + x, y = _inputs("multiclass") + with pytest.raises(ValueError, match="x_is_test_only=True requires kv_cache"): + arch(x[NUM_TRAIN:], y, task_type="multiclass", x_is_test_only=True) + + +@torch.no_grad() +def test__kv_cache__row_chunked_matches_unchunked() -> None: + """Cached forward with a small inference_row_chunk_size must match unchunked.""" + arch = _get_model() + x, y = _inputs("regression") + perf = PerformanceOptions(use_chunkwise_inference=False) + + out_standard = arch(x, y, task_type="regression", performance_options=perf) + _, cache = arch( + x, y, task_type="regression", performance_options=perf, return_kv_cache=True + ) + + # Force multi-chunk test-row processing: 4 test rows / 3 per chunk = 2 chunks. + arch.inference_row_chunk_size = 3 + out_cached_chunked = arch( + x, y, task_type="regression", performance_options=perf, kv_cache=cache + ) + torch.testing.assert_close(out_cached_chunked, out_standard, rtol=0, atol=1e-5) + + +@torch.no_grad() +@pytest.mark.parametrize( + "config_overrides", + [ + {"icl_num_kv_heads": 2, "icl_num_kv_heads_test": 1}, + {"icl_num_kv_heads": 4, "icl_num_kv_heads_test": 2}, + {"icl_num_kv_heads_test": None}, + ], +) +def test__kv_cache__gqa_variants_match_standard( + config_overrides: dict[str, object], +) -> None: + """KV-cache inference with GQA / MQA head layouts reproduces the forward.""" + arch = _get_model(**config_overrides) + x, y = _inputs("regression") + + out_standard = arch(x, y, task_type="regression") + _, cache = arch(x, y, task_type="regression", return_kv_cache=True) + out_cached = arch( + x[NUM_TRAIN:], y, task_type="regression", kv_cache=cache, x_is_test_only=True + ) + torch.testing.assert_close(out_cached, out_standard, rtol=0, atol=1e-5) + + +@torch.no_grad() +def test__kv_cache__regression_caches_no_decoder_keys() -> None: + """Regression has no many-class decoder, so its cache omits that term.""" + arch = _get_model() + x, y = _inputs("regression") + _, cache = arch(x, y, task_type="regression", return_kv_cache=True) + assert cache.decoder_keys is None + + x_cls, y_cls = _inputs("multiclass") + _, cls_cache = arch(x_cls, y_cls, task_type="multiclass", return_kv_cache=True) + decoder = arch.heads.many_class_decoder + assert cls_cache.decoder_keys.shape == ( + BATCH, + NUM_TRAIN, + decoder.num_heads, + decoder.head_dim, + ) + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__kv_cache__cached_path__omits_train_embeddings(task_type: TaskType) -> None: + """The cached path cannot report train embeddings; only the keys survive.""" + arch = _get_model() + x, y = _inputs(task_type) + _, cache = arch(x, y, task_type=task_type, return_kv_cache=True) + output = arch( + x[NUM_TRAIN:], + y, + task_type=task_type, + kv_cache=cache, + x_is_test_only=True, + only_return_standard_out=False, + ) + assert set(output) == {"standard", "test_embeddings"} + + +@torch.no_grad() +@pytest.mark.parametrize("use_chunkwise", [False, True]) +@pytest.mark.parametrize( + "autocast_dtype", + [ + torch.float16, + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif( + sys.platform == "win32" and not torch.cuda.is_available(), + reason=( + "bf16 CPU kernels crash with STATUS_ILLEGAL_INSTRUCTION " + "(0xc000001d) on Windows CI runners" + ), + ), + ), + ], +) +def test__kv_cache__works_under_autocast( + use_chunkwise: bool, autocast_dtype: torch.dtype +) -> None: + """An fp32 cache is usable under an fp16/bf16 autocast forward.""" + arch = _get_model() + x, y = _inputs("regression") + perf = PerformanceOptions(use_chunkwise_inference=use_chunkwise) + + out_standard = arch(x, y, task_type="regression", performance_options=perf) + _, cache = arch( + x, y, task_type="regression", performance_options=perf, return_kv_cache=True + ) + + device_type = "cuda" if torch.cuda.is_available() else "cpu" + with torch.autocast(device_type=device_type, dtype=autocast_dtype): + out_cached_autocast = arch( + x, y, task_type="regression", performance_options=perf, kv_cache=cache + ) + out_standard_autocast = arch( + x, y, task_type="regression", performance_options=perf + ) + + # Autocast introduces precision differences; use a loose tolerance. bf16 + # carries three significant digits, so it needs the looser one. + atol = 2e-2 if autocast_dtype == torch.bfloat16 else 1e-2 + torch.testing.assert_close( + out_cached_autocast.float(), out_standard.float(), rtol=0, atol=atol + ) + torch.testing.assert_close( + out_cached_autocast.float(), out_standard_autocast.float(), rtol=0, atol=atol + ) + + +# --------------------------------------------------------------------------- +# KV cache quantization +# --------------------------------------------------------------------------- + + +@torch.no_grad() +@pytest.mark.parametrize("dtype", [QUANTIZED_KV_DTYPE, FP8_KV_DTYPE]) +def test__quantize__kv_entries__only_the_kv_is_converted(dtype: torch.dtype) -> None: + arch = _get_model() + x, y = _inputs("multiclass") + _, cache = arch(x, y, task_type="multiclass", return_kv_cache=True) + quantized = cache.quantize(dtype) + + assert all(isinstance(e, QuantizedKVCacheEntry) for e in quantized.kv.values()) + assert all(e.key.dtype == dtype for e in quantized.kv.values()) + # The ECDF context is already narrow; the KV quantization leaves it alone. + assert cache.ecdf_context.dtype == tabpfn_v3_5.ECDF_CONTEXT_DTYPE + assert quantized.ecdf_context is cache.ecdf_context + assert set(cache.scaler_cache) == {"mean", "std"} + # Everything outside the KV cache is passed through untouched. + assert quantized.decoder_keys is cache.decoder_keys + assert quantized.scaler_cache is cache.scaler_cache + assert quantized.inducing_hidden is cache.inducing_hidden + assert quantized.train_shape == cache.train_shape + + +def test__kv_cache__quantize_passthrough_on_already_quantized() -> None: + """quantize() must not re-quantize existing QuantizedKVCacheEntry values.""" + torch.manual_seed(0) + entry = KVCacheEntry(key=torch.randn(1, 4, 1, 2), value=torch.randn(1, 4, 1, 2)) + cache = TabPFNV3p5Cache(kv={0: entry}) + q1 = cache.quantize() + q2 = q1.quantize() + assert isinstance(q2.kv[0], QuantizedKVCacheEntry) + assert q1.kv[0] is q2.kv[0] + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", ["multiclass", "regression"]) +def test__forward__kv_cache__missing_cells_in_train_and_test__matches_uncached( + task_type: TaskType, +) -> None: + """Filled test cells must tie with the filled train cells they rank against.""" + arch = _get_model() + x, y = _inputs(task_type) + gen = torch.Generator().manual_seed(1) + x = x.masked_fill(torch.rand(x.shape, generator=gen) < 0.4, float("nan")) + full = arch(x, y, task_type=task_type) + _, cache = arch(x, y, task_type=task_type, return_kv_cache=True) + from_cache = arch( + x[NUM_TRAIN:], y, task_type=task_type, kv_cache=cache, x_is_test_only=True + ) + torch.testing.assert_close(full, from_cache, rtol=0, atol=1e-5) + + +@torch.no_grad() +@pytest.mark.parametrize("cache_dtype", [QUANTIZED_KV_DTYPE, FP8_KV_DTYPE]) +def test__kv_cache__layerwise_quantization_matches_post_forward( + cache_dtype: torch.dtype, +) -> None: + """Quantizing during construction produces the same cache as afterward.""" + arch = _get_model() + x, y = _inputs("multiclass") + + _, full_precision = arch(x, y, task_type="multiclass", return_kv_cache=True) + _, layerwise = arch( + x, + y, + task_type="multiclass", + return_kv_cache=True, + performance_options=PerformanceOptions(kv_cache_dtype=cache_dtype), + ) + post_forward = full_precision.quantize(cache_dtype) + + assert layerwise.decoder_keys.dtype == full_precision.decoder_keys.dtype + for layer_idx in post_forward.kv: + expected = post_forward.kv[layer_idx] + actual = layerwise.kv[layer_idx] + assert isinstance(expected, QuantizedKVCacheEntry) + assert isinstance(actual, QuantizedKVCacheEntry) + # torch.equal lacks CPU float8 support in the lowest supported PyTorch. + # Comparing after an exact float32 widening works for int8 and float8. + assert torch.equal(actual.key.float(), expected.key.float()) + assert torch.equal(actual.value.float(), expected.value.float()) + # The scales are an absmax over a fresh forward pass, and BLAS on some + # platforms (macOS arm64) is not bitwise reproducible across runs. + torch.testing.assert_close( + actual.key_scale, expected.key_scale, rtol=1e-6, atol=0 + ) + torch.testing.assert_close( + actual.value_scale, expected.value_scale, rtol=1e-6, atol=0 + ) + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +def test__forward__quantized_kv_cache__equals_the_dequantized_cache( + task_type: TaskType, +) -> None: + """Quantizing must add nothing but the dequantize on the way back in.""" + arch = _get_model() + x, y = _inputs(task_type) + _, cache = arch(x, y, task_type=task_type, return_kv_cache=True) + quantized = cache.quantize() + dequantized = dataclasses.replace( + quantized, + kv={i: e.dequantize(torch.float32) for i, e in quantized.kv.items()}, + ) + predict = functools.partial( + arch, x[NUM_TRAIN:], y, task_type=task_type, x_is_test_only=True + ) + assert torch.equal(predict(kv_cache=quantized), predict(kv_cache=dequantized)) + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +@pytest.mark.parametrize("use_chunkwise", [False, True]) +def test__quantized_kv_cache__close_to_standard_forward( + task_type: TaskType, use_chunkwise: bool +) -> None: + """Int8-quantized KV cache produces output close to the standard forward. + + Decomposes error so a regression in the cache path itself (which should match + standard at near machine precision) can't hide behind the loose int8 tolerance. + """ + arch = _get_model() + x, y = _inputs(task_type) + perf = PerformanceOptions(use_chunkwise_inference=use_chunkwise) + + out_standard = arch(x, y, task_type=task_type, performance_options=perf) + _, cache = arch( + x, y, task_type=task_type, performance_options=perf, return_kv_cache=True + ) + out_cached = arch( + x, y, task_type=task_type, performance_options=perf, kv_cache=cache + ) + out_quantized = arch( + x, y, task_type=task_type, performance_options=perf, kv_cache=cache.quantize() + ) + + torch.testing.assert_close(out_cached, out_standard, rtol=0, atol=1e-5) + torch.testing.assert_close(out_quantized, out_cached, rtol=0, atol=1e-2) + + +# --------------------------------------------------------------------------- +# Cache size +# --------------------------------------------------------------------------- + + +def _sum_cache_tensors(obj: object) -> int: + """Recursively sum ``numel * element_size`` over every tensor in a cache. + + Walks dataclasses / dicts / lists so a newly-added cached tensor field is + automatically included -- the completeness guard for ``get_cache_size``. + """ + if isinstance(obj, torch.Tensor): + return obj.numel() * obj.element_size() + if isinstance(obj, dict): + return sum(_sum_cache_tensors(v) for v in obj.values()) + if isinstance(obj, (list, tuple)): + return sum(_sum_cache_tensors(v) for v in obj) + if dataclasses.is_dataclass(obj): + return sum( + _sum_cache_tensors(getattr(obj, f.name)) for f in dataclasses.fields(obj) + ) + return 0 + + +def _quantize( + cache: TabPFNV3p5Cache, kv_cache_precision: Literal["auto", "int8", "fp8"] +) -> TabPFNV3p5Cache: + """Apply the quantization step the inference engine applies.""" + if kv_cache_precision == "int8": + return cache.quantize() + if kv_cache_precision == "fp8": + return cache.quantize(FP8_KV_DTYPE) + return cache + + +@torch.no_grad() +@pytest.mark.parametrize("task_type", TASK_TYPES) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("kv_cache_precision", ["int8", "fp8", "auto"]) +@pytest.mark.parametrize("config_overrides", [{}, {"icl_num_kv_heads_test": None}]) +def test__get_cache_size__matches_whole_cache( + task_type: TaskType, + kv_cache_precision: Literal["auto", "int8", "fp8"], + dtype: torch.dtype, + config_overrides: dict[str, object], +) -> None: + """get_cache_size equals the exact byte size of every tensor in a real cache. + + Parametrized over ``dtype`` to cover the engine's forced-precision path: the + engine casts the model and inputs to that dtype and runs the forward with + autocast disabled, so every non-KV term lands at that one dtype. + """ + arch = _get_model(**config_overrides) + arch.type(dtype) # mirror the engine's set_dtype for forced precision. + # get_cache_size describes one estimator, so batch size 1. + x, y = _inputs(task_type, batch=1) + x, y = x.to(dtype), y.to(dtype) + _, cache = arch(x, y, task_type=task_type, return_kv_cache=True) + cache = _quantize(cache, kv_cache_precision) + + total = get_cache_size( + n_train=NUM_TRAIN, + n_features=NUM_FEATURES, + model_config=_config(config_overrides), + task_type=task_type, + base_dtype=dtype, + kv_cache_precision=kv_cache_precision, + ) + assert total == _sum_cache_tensors(cache) + + +@torch.no_grad() +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Autocast inference is only enabled on CUDA (disabled on CPU/MPS), so " + "the mixed-precision cache it produces can only be built on a GPU.", +) +@pytest.mark.parametrize("task_type", TASK_TYPES) +@pytest.mark.parametrize("kv_cache_precision", ["int8", "fp8", "auto"]) +def test__get_cache_size__matches_whole_cache_autocast( + task_type: TaskType, kv_cache_precision: Literal["auto", "int8", "fp8"] +) -> None: + """get_cache_size matches a real cache built on the GPU autocast path. + + Autocast keeps fp32 model weights and casts ops at runtime, so the cache mixes + dtypes; ``get_cache_size(base_dtype="autocast")`` must size each term at its + real precision and still match to the byte. + """ + device = torch.device("cuda") + arch = _get_model().to(device) + x, y = _inputs(task_type, batch=1) + x, y = x.to(device), y.to(device) + with get_autocast_context(device, enabled=True): + _, cache = arch(x, y, task_type=task_type, return_kv_cache=True) + cache = _quantize(cache, kv_cache_precision) + + total = get_cache_size( + n_train=NUM_TRAIN, + n_features=NUM_FEATURES, + model_config=_config({}), + task_type=task_type, + base_dtype="autocast", + kv_cache_precision=kv_cache_precision, + ) + assert total == _sum_cache_tensors(cache) + + +def test__get_cache_size__mqa_smaller_than_mha() -> None: + """Fewer cached KV heads (MQA on the test partition) shrinks the KV term.""" + common = {**_SMALL_CONFIG, "icl_num_kv_heads_test": None} + mha, _ = tabpfn_v3_5.parse_config(common) # H_kv = icl_num_heads = 4 + mqa, _ = tabpfn_v3_5.parse_config({**common, "icl_num_kv_heads_test": 1}) + n_train = 50 + # kv_cache_precision defaults to "int8", so the KV cache is int8 (1 byte) + # regardless of base_dtype; base_dtype only sizes the (cancelling) non-KV terms. + kw = { + "n_train": n_train, + "n_features": 5, + "task_type": "multiclass", + "base_dtype": torch.float32, + } + est_mha = get_cache_size(model_config=mha, **kw) + est_mqa = get_cache_size(model_config=mqa, **kw) + + # mha and mqa differ ONLY in the KV term (H_kv 4 vs 1); every other term + # (activations, inducing, scaler, ecdf) is identical, so it cancels in the diff. + icl_emsize = mha.embed_dim * mha.feat_agg_num_cls_tokens + head_dim = icl_emsize // mha.icl_num_heads + kv_per_head = mha.nlayers * 2 * n_train * head_dim # int8 KV -> 1 byte/element + assert est_mqa < est_mha + assert est_mha - est_mqa == (4 - 1) * kv_per_head + + +def test__get_cache_size__regression_omits_the_decoder_keys() -> None: + config, _ = tabpfn_v3_5.parse_config(_SMALL_CONFIG) + kw = { + "n_train": 50, + "n_features": 5, + "model_config": config, + "base_dtype": torch.float32, + } + multiclass = get_cache_size(task_type="multiclass", **kw) + regression = get_cache_size(task_type="regression", **kw) + decoder_keys = 50 * config.decoder_num_heads * config.decoder_head_dim * 4 + assert multiclass - regression == decoder_keys + + +def test__get_cache_size__ecdf_term_stops_growing_at_the_bucket_count() -> None: + """Past `cell_ecdf_num_buckets` train rows, only the KV cache keeps growing.""" + config, _ = tabpfn_v3_5.parse_config({**_SMALL_CONFIG, "cell_ecdf_num_buckets": 16}) + kw = { + "n_features": 5, + "model_config": config, + "task_type": "regression", + "base_dtype": torch.float32, + "kv_cache_precision": "auto", + } + head_dim = config.embed_dim * config.feat_agg_num_cls_tokens // config.icl_num_heads + kv_per_row = config.nlayers * 2 * config.icl_num_kv_heads_test * head_dim * 4 + below = get_cache_size(n_train=16, **kw) + above = get_cache_size(n_train=32, **kw) + assert above - below == 16 * kv_per_row + + +def test__get_cache_size__invalid_precision__raises() -> None: + config, _ = tabpfn_v3_5.parse_config(_SMALL_CONFIG) + with pytest.raises(ValueError, match="Invalid kv_cache_precision"): + get_cache_size( + n_train=10, + n_features=5, + model_config=config, + task_type="multiclass", + base_dtype=torch.float32, + kv_cache_precision="int4", # type: ignore[arg-type] + ) + + +@pytest.mark.slow +def test__get_cache_size__tabpfn3_5_classifier_1000_rows() -> None: + """Pin get_cache_size for the real TabPFN-v3.5 checkpoint at 1,000 train rows + (1 estimator, engine defaults: int8 KV, fp16 rest). + """ + clf = TabPFNClassifier.create_default_for_version(ModelVersion.V3_5) + # Loads the checkpoint (config + weights) without needing fit data. + clf._initialize_model_variables() + config = clf.configs_[0] + assert isinstance(config, TabPFNV3p5Config) + + n_train, n_features = 1000, 1 + total = get_cache_size( + n_train=n_train, + n_features=n_features, + model_config=config, + task_type="multiclass", + base_dtype=torch.float16, + kv_cache_precision="int8", + ) + + # Fixed terms for the shipped config (nlayers 24, H_kv 1, head_dim 64, 6x64 + # decoder), all at fp16 apart from the int8 KV: + # KV int8: 24 * 2 * 1 * 64 * 1000 = 3,072,000 + # + int8 KV scales: 24 * 2 * 2 bytes = 96 + # + scaler stats: 2 * n_features * 2 bytes = 4 + # + fp16 decoder keys: 6 * 64 * 1000 * 2 = 768,000 + # The inducing and ECDF terms depend on the shipped embedder config, so they + # are derived from it instead of hardcoded. + inducing = ( + config.dist_embed_num_blocks + * n_features + * config.dist_embed_num_inducing_points + * config.embed_dim + ) * torch.float16.itemsize + ecdf = ( + 3 + * n_features + * min(config.cell_ecdf_num_buckets, n_train) + * tabpfn_v3_5.ECDF_CONTEXT_DTYPE.itemsize + ) + # Numbers need manual update if we bump the default architecture. + assert total == 3_072_000 + 96 + 4 + 768_000 + inducing + ecdf diff --git a/tests/test_browser_auth.py b/tests/test_browser_auth.py index cf49f19ab..e874b3cc6 100644 --- a/tests/test_browser_auth.py +++ b/tests/test_browser_auth.py @@ -15,13 +15,18 @@ import pytest from tabpfn.browser_auth import ( + _get_license_name, _has_display, delete_cached_token, get_cached_token, save_token, verify_token, ) -from tabpfn.errors import TabPFNLicenseError +from tabpfn.errors import ( + TabPFNError, + TabPFNHuggingFaceGatedRepoError, + TabPFNLicenseError, +) # --------------------------------------------------------------------------- # Fixtures @@ -639,3 +644,30 @@ def test_graphical_opens_browser(self): assert result == "browser-jwt" mock_browser.assert_called_once() + + +def _http_error(code: int) -> urllib.error.HTTPError: + return urllib.error.HTTPError("https://huggingface.co", code, "", {}, None) # type: ignore[arg-type] + + +@pytest.mark.parametrize("code", [401, 404]) +def test__get_license_name__unreleased_repo_not_visible__says_not_available_yet( + code: int, +) -> None: + with ( + patch( + "tabpfn.browser_auth.urllib.request.urlopen", side_effect=_http_error(code) + ), + pytest.raises(TabPFNError, match="not publicly available yet"), + ): + _get_license_name("tabpfn_3_5") + + +def test__get_license_name__released_repo_not_visible__raises_gated_error() -> None: + with ( + patch( + "tabpfn.browser_auth.urllib.request.urlopen", side_effect=_http_error(401) + ), + pytest.raises(TabPFNHuggingFaceGatedRepoError), + ): + _get_license_name("tabpfn_3") diff --git a/tests/test_classifier_interface.py b/tests/test_classifier_interface.py index 0b8e200f3..4c1c0767b 100644 --- a/tests/test_classifier_interface.py +++ b/tests/test_classifier_interface.py @@ -67,6 +67,8 @@ def X_y() -> tuple[np.ndarray, np.ndarray]: ModelSource.get_classifier_v2(), ModelSource.get_classifier_v2_5(), ModelSource.get_classifier_v3(), + ModelSource.get_v3_5(), + ModelSource.get_v3_5_fast(), ] fit_modes = ["low_memory", "fit_preprocessors"] @@ -530,7 +532,14 @@ def test_balance_probabilities_alters_proba_output() -> None: @pytest.mark.parametrize( "model_version", - [ModelVersion.V2, ModelVersion.V2_5, ModelVersion.V2_6, ModelVersion.V3], + [ + ModelVersion.V2, + ModelVersion.V2_5, + ModelVersion.V2_6, + ModelVersion.V3, + ModelVersion.V3_5, + ModelVersion.V3_5_FAST, + ], ) # Disable MPS as it doesn't support float64. @pytest.mark.parametrize("device", [d for d in get_pytest_devices() if d != "mps"]) @@ -565,12 +574,15 @@ def test__fit_preprocessors_and_with_cache_produce_equal_results( np.testing.assert_array_equal(preds, tabpfn.predict(X)) +@pytest.mark.parametrize( + "model_version", [ModelVersion.V3, ModelVersion.V3_5, ModelVersion.V3_5_FAST] +) @pytest.mark.parametrize("device", get_pytest_devices()) -def test__fit_preprocessors_and_with_cache_with_quantized_kv_cache__v3( - X_y: tuple[np.ndarray, np.ndarray], device: str +def test__fit_preprocessors_and_with_cache_with_quantized_kv_cache__v3_family( + X_y: tuple[np.ndarray, np.ndarray], model_version: ModelVersion, device: str ) -> None: kwargs = { - "version": ModelVersion.V3, + "version": model_version, "n_estimators": 2, "inference_precision": torch.float32, "random_state": 0, @@ -1442,6 +1454,34 @@ def test__create_default_for_version__v3__uses_correct_defaults() -> None: assert "-v3-" in estimator.model_path +def test__create_default_for_version__v3_5__uses_correct_defaults() -> None: + estimator = TabPFNClassifier.create_default_for_version(ModelVersion.V3_5) + + assert isinstance(estimator, TabPFNClassifier) + assert estimator.n_estimators == "auto" + assert estimator.softmax_temperature == "auto" + assert isinstance(estimator.model_path, str) + # One multitask checkpoint backs both estimators, so the estimator type is not + # part of the file name. + assert "-v3.5-" in estimator.model_path + assert "fast" not in estimator.model_path + + +def test__create_default_for_version__v3_5_fast__uses_correct_defaults() -> None: + estimator = TabPFNClassifier.create_default_for_version(ModelVersion.V3_5_FAST) + + assert isinstance(estimator, TabPFNClassifier) + assert estimator.n_estimators == "auto" + assert estimator.softmax_temperature == "auto" + assert isinstance(estimator.model_path, str) + assert "-v3.5-fast-" in estimator.model_path + + +def test__create_default_for_version__unknown_version__raises() -> None: + with pytest.raises(ValueError, match="Unknown version"): + TabPFNClassifier.create_default_for_version("v99") # type: ignore[arg-type] + + def test__create_default_for_version__passes_through_overrides() -> None: estimator = TabPFNClassifier.create_default_for_version( ModelVersion.V2_5, n_estimators=16 diff --git a/tests/test_inference_config.py b/tests/test_inference_config.py index 3f83d16e5..49a6cf9e9 100644 --- a/tests/test_inference_config.py +++ b/tests/test_inference_config.py @@ -199,8 +199,13 @@ def test__regressor_get_inference_config__with_override__applies_override() -> N assert specs.inference_config.POLYNOMIAL_FEATURES == "no" -def test__cpu_sample_limit__v3__returns_5000() -> None: - assert cpu_sample_limit(ModelVersion.V3) == 5000 +@pytest.mark.parametrize( + "model_version", [ModelVersion.V3, ModelVersion.V3_5, ModelVersion.V3_5_FAST] +) +def test__cpu_sample_limit__v3_onwards__returns_5000( + model_version: ModelVersion, +) -> None: + assert cpu_sample_limit(model_version) == 5000 def test__cpu_sample_limit__pre_v3_versions__return_1000() -> None: diff --git a/tests/test_model_loading.py b/tests/test_model_loading.py index e94d0dd87..88900c1f2 100644 --- a/tests/test_model_loading.py +++ b/tests/test_model_loading.py @@ -18,7 +18,7 @@ from torch import Tensor, nn from tabpfn import model_loading -from tabpfn.architectures import ARCHITECTURES, tabpfn_v2, tabpfn_v3 +from tabpfn.architectures import ARCHITECTURES, tabpfn_v2, tabpfn_v3, tabpfn_v3_5 from tabpfn.architectures.interface import ( Architecture, ArchitectureConfig, @@ -26,6 +26,8 @@ ) from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution from tabpfn.architectures.tabpfn_v3 import TabPFNV3Config +from tabpfn.architectures.tabpfn_v3_5 import TabPFNV3p5Config +from tabpfn.checkpoint import save_as_safetensors from tabpfn.constants import ModelVersion from tabpfn.inference_config import DEFAULT_SOFTMAX_TEMPERATURE, InferenceConfig from tabpfn.preprocessing import PreprocessorConfig @@ -167,6 +169,26 @@ def test__save_tabpfn_model__stores_v3_architecture_and_inference_config( assert checkpoint["inference_config"] == asdict(inference_config) +def test__save_tabpfn_model__stores_v3_5_architecture_and_inference_config( + tmp_path: Path, +) -> None: + config = TabPFNV3p5Config(max_num_classes=10, num_buckets=100) + inference_config = InferenceConfig.get_default("multiclass", ModelVersion.V2_5) + estimator = SimpleNamespace( + models_=[torch.nn.Linear(1, 1)], + configs_=[config], + inference_config_=inference_config, + ) + checkpoint_path = tmp_path / "checkpoint.ckpt" + + model_loading.save_tabpfn_model(estimator, checkpoint_path) + + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + assert checkpoint["architecture_name"] == "tabpfn_v3_5" + assert checkpoint["config"]["name"] == "TabPFN-v3.5" + assert checkpoint["inference_config"] == asdict(inference_config) + + def test__load_v2_checkpoint__returns_v2_preprocessings( tmp_path: Path, ) -> None: @@ -333,6 +355,81 @@ def _build_small_v3_checkpoint( } +def _build_small_v3_5_checkpoint( + inference_config: InferenceConfig, + *, + max_num_classes: int, +) -> dict: + config = TabPFNV3p5Config( + max_num_classes=max_num_classes, + num_buckets=5, + embed_dim=48, + nlayers=1, + icl_num_heads=3, + dist_embed_num_heads=3, + feat_agg_num_heads=3, + ) + model = tabpfn_v3_5.get_architecture(config, cache_trainset_representation=False) + return { + "state_dict": model.state_dict(), + "config": asdict(config), + "architecture_name": "tabpfn_v3_5", + "inference_config": asdict(inference_config), + } + + +@pytest.mark.parametrize("estimator_type", ["classifier", "regressor"]) +@pytest.mark.parametrize("version", ["v3.5", "v3.5-fast"]) +def test__load_v3_5_multitask_ckpt__backs_both_estimator_types( + tmp_path: Path, + estimator_type: Literal["classifier", "regressor"], + version: Literal["v3.5", "v3.5-fast"], +) -> None: + """A v3.5 checkpoint carries both heads and its own inference config.""" + inference_config = InferenceConfig( + PREPROCESS_TRANSFORMS=[PreprocessorConfig("quantile_uni_coarse")] + ) + checkpoint = _build_small_v3_5_checkpoint(inference_config, max_num_classes=10) + checkpoint_path = tmp_path / f"tabpfn-{version}-test.safetensors" + save_as_safetensors(checkpoint, checkpoint_path) + + models, criterion, configs, loaded_inference_config = ( + model_loading.load_model_criterion_config( + model_path=[checkpoint_path], + check_bar_distribution_criterion=estimator_type == "regressor", + cache_trainset_representation=False, + estimator_type=estimator_type, + version=version, + download_if_not_exists=False, + ) + ) + + assert isinstance(models[0], tabpfn_v3_5.TabPFNV3p5) + assert isinstance(configs[0], TabPFNV3p5Config) + if estimator_type == "regressor": + assert isinstance(criterion, FullSupportBarDistribution) + else: + assert isinstance(criterion, nn.CrossEntropyLoss) + assert loaded_inference_config == inference_config + + +@pytest.mark.parametrize( + ("file_name", "expected"), + [ + ("tabpfn-v3.5-fast-20260909.safetensors", ModelVersion.V3_5_FAST), + ("tabpfn-v3.5-20260909.safetensors", ModelVersion.V3_5), + ("tabpfn-v3-classifier-v3_default.ckpt", ModelVersion.V3), + ("tabpfn-v2.6-regressor-v2.6_default.ckpt", ModelVersion.V2_6), + ("tabpfn-v2.5-classifier-v2.5_default.ckpt", ModelVersion.V2_5), + ("tabpfn-v2-classifier.ckpt", ModelVersion.V2), + ], +) +def test__resolve_model_version__reads_the_version_off_the_file_name( + tmp_path: Path, file_name: str, expected: ModelVersion +) -> None: + assert model_loading.resolve_model_version(tmp_path / file_name) == expected + + def test__load_v3_classification_ckpt__returns_inference_config_from_checkpoint( tmp_path: Path, ) -> None: diff --git a/tests/test_regressor_interface.py b/tests/test_regressor_interface.py index 9bc16d8e7..40af51c05 100644 --- a/tests/test_regressor_interface.py +++ b/tests/test_regressor_interface.py @@ -57,6 +57,8 @@ ModelSource.get_regressor_v2(), ModelSource.get_regressor_v2_5(), ModelSource.get_regressor_v3(), + ModelSource.get_v3_5(), + ModelSource.get_v3_5_fast(), ] fit_modes = ["low_memory", "fit_preprocessors"] @@ -298,7 +300,14 @@ def test__fit_predict__specify_inference_config__outputs_correct_shape( @pytest.mark.parametrize( "model_version", - [ModelVersion.V2, ModelVersion.V2_5, ModelVersion.V2_6, ModelVersion.V3], + [ + ModelVersion.V2, + ModelVersion.V2_5, + ModelVersion.V2_6, + ModelVersion.V3, + ModelVersion.V3_5, + ModelVersion.V3_5_FAST, + ], ) # Disable MPS as it doesn't support float64. @pytest.mark.parametrize("device", [d for d in get_pytest_devices() if d != "mps"]) @@ -333,12 +342,15 @@ def test__fit_preprocessors_and_with_cache_produce_equal_results( np.testing.assert_array_almost_equal(preds, tabpfn.predict(X), decimal=2) +@pytest.mark.parametrize( + "model_version", [ModelVersion.V3, ModelVersion.V3_5, ModelVersion.V3_5_FAST] +) @pytest.mark.parametrize("device", get_pytest_devices()) -def test__fit_preprocessors_and_with_cache_with_quantized_kv_cache__v3( - X_y: tuple[np.ndarray, np.ndarray], device: str +def test__fit_preprocessors_and_with_cache_with_quantized_kv_cache__v3_family( + X_y: tuple[np.ndarray, np.ndarray], model_version: ModelVersion, device: str ) -> None: kwargs = { - "version": ModelVersion.V3, + "version": model_version, "n_estimators": 2, "inference_precision": torch.float32, "random_state": 0, @@ -1019,6 +1031,34 @@ def test__create_default_for_version__v3__uses_correct_defaults() -> None: assert "-v3-" in estimator.model_path +def test__create_default_for_version__v3_5__uses_correct_defaults() -> None: + estimator = TabPFNRegressor.create_default_for_version(ModelVersion.V3_5) + + assert isinstance(estimator, TabPFNRegressor) + assert estimator.n_estimators == "auto" + assert estimator.softmax_temperature == "auto" + assert isinstance(estimator.model_path, str) + # One multitask checkpoint backs both estimators, so the estimator type is not + # part of the file name. + assert "-v3.5-" in estimator.model_path + assert "fast" not in estimator.model_path + + +def test__create_default_for_version__v3_5_fast__uses_correct_defaults() -> None: + estimator = TabPFNRegressor.create_default_for_version(ModelVersion.V3_5_FAST) + + assert isinstance(estimator, TabPFNRegressor) + assert estimator.n_estimators == "auto" + assert estimator.softmax_temperature == "auto" + assert isinstance(estimator.model_path, str) + assert "-v3.5-fast-" in estimator.model_path + + +def test__create_default_for_version__unknown_version__raises() -> None: + with pytest.raises(ValueError, match="Unknown version"): + TabPFNRegressor.create_default_for_version("v99") # type: ignore[arg-type] + + def test__create_default_for_version__passes_through_overrides() -> None: estimator = TabPFNRegressor.create_default_for_version( ModelVersion.V2_5, n_estimators=16