zsfm_tabpfn/infer/
mod.rs

1//! TabPFN-3 inference engine — classification only, single pass (no ensembling). NON-COMMERCIAL
2//! WEIGHTS LICENSE (TabPFN-3 Non-Commercial License v1.0): research/internal/benchmarking use
3//! only — no production, commercial, or hosted-service use without a separate license from Prior
4//! Labs GmbH. See the crate's `Cargo.toml` description.
5//!
6//! Architecture (three stacked transformers, closely related to TabICL's but with real
7//! differences — RMSNorm throughout instead of LayerNorm, unbiased separate Q/K/V projections
8//! instead of a packed `in_proj`, no-bias MLPs, and a very different final decoder):
9//!
10//! 1. **Feature distribution embedder**: each (grouped, NaN-indicator-augmented) feature column
11//!    is embedded independently by a shared 3-block Set Transformer (`InducedSelfAttentionBlock`:
12//!    learned inducing points cross-attend to the *training* rows only — with a learned
13//!    query-aware elementwise attention scale ("SoftmaxScalingMLP", identical in spirit to
14//!    TabICL's SSMax) — then the full column cross-attends back to the refined inducing points),
15//!    with the training targets folded in beforehand via an orthogonal class-embedding lookup.
16//! 2. **Column aggregator**: per row, the (grouped) feature-column embeddings plus 4 learned CLS
17//!    tokens attend to each other (non-interleaved RoPE, frequencies stored in the checkpoint);
18//!    the final block reads out via CLS-tokens-as-query cross-attention, concatenated into one
19//!    `embed_dim * 4` row representation.
20//! 3. **ICL transformer**: training targets are folded into their rows' representations (again
21//!    via an orthogonal embedding lookup), then a 24-block transformer (SoftmaxScalingMLP on
22//!    every block) lets query rows attend to training rows only — with a GQA-style quirk: query
23//!    (test) rows attend using only the *first* of the 8 K/V heads (broadcast to all 8 query
24//!    heads), while training rows use the full 8 K/V heads. A `ManyClassDecoder` then reads out
25//!    a probability-like distribution per class via one more attention pass — Q/K project the
26//!    row embeddings, V is the (per-head-broadcast) one-hot training-label encoding, so the
27//!    attention output is literally an attention-weighted average of one-hot labels; the result
28//!    is log-transformed into logits.
29//!
30//! Regression (the bar-distribution/quantile head) is out of scope — see the crate's `Cargo.toml`
31//! description. NaN/Inf indicator features are always computed (the checkpoint's `x_embed`
32//! expects them) but real missing-value handling is not exercised: this port assumes clean
33//! (non-NaN) input, matching the scope decision already established for every other model in
34//! this workspace.
35
36use std::io::{BufReader, Read, Seek};
37use std::path::Path;
38
39use anyhow::{Context, Result};
40use candle_core::quantized::gguf_file;
41use candle_core::{DType, Device, Tensor};
42
43use crate::config::TabPfnConfig;
44
45/// `nn.RMSNorm`'s default `eps` when unset: the input dtype's machine epsilon. This checkpoint's
46/// compute dtype is F32, so `torch.finfo(torch.float32).eps`.
47const RMS_EPS: f64 = 1.192_092_9e-7;
48
49// ---------------------------------------------------------------------------
50// Weight structs
51// ---------------------------------------------------------------------------
52
53struct SsmaxW {
54    base_fc1_w: Tensor,
55    base_fc1_b: Tensor,
56    base_fc2_w: Tensor,
57    base_fc2_b: Tensor,
58    query_fc1_w: Tensor,
59    query_fc1_b: Tensor,
60    query_fc2_w: Tensor,
61    query_fc2_b: Tensor,
62}
63
64/// Unbiased separate Q/K/V/out projections (`nn.Linear(..., bias=False)`), shared shape for
65/// `Attention`, `CrossAttention`, and `ICLAttention`.
66struct QkvW {
67    q_w: Tensor,
68    k_w: Tensor,
69    v_w: Tensor,
70    out_w: Tensor,
71    ssmax: Option<SsmaxW>,
72}
73
74/// No-bias two-layer GELU MLP (`nn.Sequential(Linear, GELU, Linear)`, both `bias=False`).
75struct MlpW {
76    fc1_w: Tensor,
77    fc2_w: Tensor,
78}
79
80/// `CrossAttentionBlock`: pre-norm cross-attention with *separate* Q-side/KV-side RMSNorms
81/// (not the same weights reused, unlike TabICL's `MultiheadAttentionBlock`).
82struct CrossAttnBlockW {
83    attn: QkvW,
84    mlp: MlpW,
85    ln_q_w: Tensor,
86    ln_kv_w: Tensor,
87    ln2_w: Tensor,
88}
89
90struct IsabW {
91    ind_vectors: Tensor,
92    block1: CrossAttnBlockW,
93    block2: CrossAttnBlockW,
94}
95
96/// `TransformerBlock`: pre-norm self-attention (`ColumnAggregator`'s blocks). A single
97/// `layernorm` is reused for both query and context sides in the CLS-readout (`forward_cross`)
98/// variant, matching the reference's literal `self.layernorm(...)` reuse.
99struct TransformerBlockW {
100    attn: QkvW,
101    mlp: MlpW,
102    ln_w: Tensor,
103    ln_mlp_w: Tensor,
104}
105
106/// `ICLTransformerBlock`: pre-norm `ICLAttention` (train-only K/V, GQA-test quirk) + MLP.
107struct IclBlockW {
108    attn: QkvW,
109    mlp: MlpW,
110    ln_w: Tensor,
111    ln_mlp_w: Tensor,
112}
113
114struct ManyClassDecoderW {
115    q_w: Tensor,
116    q_b: Tensor,
117    k_w: Tensor,
118    k_b: Tensor,
119    ssmax: Option<SsmaxW>,
120}
121
122pub struct TabPfnModel {
123    device: Device,
124    config: TabPfnConfig,
125    x_embed_w: Tensor,
126    x_embed_b: Tensor,
127    col_y_encoder_w: Tensor,
128    icl_y_encoder_w: Tensor,
129    dist_embed: Vec<IsabW>,
130    col_agg_blocks: Vec<TransformerBlockW>,
131    col_agg_cls_tokens: Tensor,
132    col_agg_rope_freqs: Vec<f32>,
133    col_agg_out_ln_w: Tensor,
134    icl_blocks: Vec<IclBlockW>,
135    output_norm_w: Tensor,
136    many_class_decoder: ManyClassDecoderW,
137}
138
139// ---------------------------------------------------------------------------
140// GGUF loading (original, dotted PyTorch tensor names — converted via the *generic*
141// `zsfm convert` path, names passed through unchanged; no crate-specific tensor_map/convert).
142// ---------------------------------------------------------------------------
143
144fn load_t(content: &gguf_file::Content, reader: &mut (impl Read + Seek), name: &str, device: &Device) -> Result<Tensor> {
145    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
146}
147
148fn load_ssmax(content: &gguf_file::Content, reader: &mut (impl Read + Seek), prefix: &str, device: &Device) -> Result<SsmaxW> {
149    let mut t = |s: &str| load_t(content, reader, &format!("{prefix}.{s}"), device);
150    Ok(SsmaxW {
151        base_fc1_w: t("base_mlp.0.weight")?,
152        base_fc1_b: t("base_mlp.0.bias")?,
153        base_fc2_w: t("base_mlp.2.weight")?,
154        base_fc2_b: t("base_mlp.2.bias")?,
155        query_fc1_w: t("query_mlp.0.weight")?,
156        query_fc1_b: t("query_mlp.0.bias")?,
157        query_fc2_w: t("query_mlp.2.weight")?,
158        query_fc2_b: t("query_mlp.2.bias")?,
159    })
160}
161
162fn load_qkv(
163    content: &gguf_file::Content,
164    reader: &mut (impl Read + Seek),
165    prefix: &str,
166    device: &Device,
167    ssmax_prefix: Option<&str>,
168) -> Result<QkvW> {
169    let mut t = |s: &str| load_t(content, reader, &format!("{prefix}.{s}"), device);
170    let q_w = t("q_projection.weight")?;
171    let k_w = t("k_projection.weight")?;
172    let v_w = t("v_projection.weight")?;
173    let out_w = t("out_projection.weight")?;
174    let ssmax = match ssmax_prefix {
175        Some(p) => Some(load_ssmax(content, reader, p, device)?),
176        None => None,
177    };
178    Ok(QkvW { q_w, k_w, v_w, out_w, ssmax })
179}
180
181fn load_mlp(content: &gguf_file::Content, reader: &mut (impl Read + Seek), prefix: &str, device: &Device) -> Result<MlpW> {
182    Ok(MlpW {
183        fc1_w: load_t(content, reader, &format!("{prefix}.0.weight"), device)?,
184        fc2_w: load_t(content, reader, &format!("{prefix}.2.weight"), device)?,
185    })
186}
187
188fn load_cross_attn_block(
189    content: &gguf_file::Content,
190    reader: &mut (impl Read + Seek),
191    prefix: &str,
192    device: &Device,
193    has_ssmax: bool,
194) -> Result<CrossAttnBlockW> {
195    let ssmax_prefix = format!("{prefix}.attn.softmax_scaling_layer");
196    let attn = load_qkv(content, reader, &format!("{prefix}.attn"), device, has_ssmax.then_some(ssmax_prefix.as_str()))?;
197    let mlp = load_mlp(content, reader, &format!("{prefix}.mlp"), device)?;
198    let ln_q_w = load_t(content, reader, &format!("{prefix}.layernorm_q.weight"), device)?;
199    let ln_kv_w = load_t(content, reader, &format!("{prefix}.layernorm_kv.weight"), device)?;
200    let ln2_w = load_t(content, reader, &format!("{prefix}.layernorm2.weight"), device)?;
201    Ok(CrossAttnBlockW { attn, mlp, ln_q_w, ln_kv_w, ln2_w })
202}
203
204fn load_transformer_block(
205    content: &gguf_file::Content,
206    reader: &mut (impl Read + Seek),
207    prefix: &str,
208    device: &Device,
209) -> Result<TransformerBlockW> {
210    let attn = load_qkv(content, reader, &format!("{prefix}.attention"), device, None)?;
211    let mlp = load_mlp(content, reader, &format!("{prefix}.mlp"), device)?;
212    let ln_w = load_t(content, reader, &format!("{prefix}.layernorm.weight"), device)?;
213    let ln_mlp_w = load_t(content, reader, &format!("{prefix}.layernorm_mlp.weight"), device)?;
214    Ok(TransformerBlockW { attn, mlp, ln_w, ln_mlp_w })
215}
216
217fn load_icl_block(
218    content: &gguf_file::Content,
219    reader: &mut (impl Read + Seek),
220    prefix: &str,
221    device: &Device,
222) -> Result<IclBlockW> {
223    let ssmax_prefix = format!("{prefix}.icl_attention.softmax_scaling_layer");
224    let attn = load_qkv(content, reader, &format!("{prefix}.icl_attention"), device, Some(&ssmax_prefix))?;
225    let mlp = load_mlp(content, reader, &format!("{prefix}.mlp"), device)?;
226    let ln_w = load_t(content, reader, &format!("{prefix}.layernorm.weight"), device)?;
227    let ln_mlp_w = load_t(content, reader, &format!("{prefix}.layernorm_mlp.weight"), device)?;
228    Ok(IclBlockW { attn, mlp, ln_w, ln_mlp_w })
229}
230
231impl TabPfnModel {
232    pub fn load(gguf_path: &Path, config: TabPfnConfig) -> Result<Self> {
233        let device = Device::Cpu;
234        let file = std::fs::File::open(gguf_path).with_context(|| format!("open {}", gguf_path.display()))?;
235        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
236        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
237
238        let x_embed_w = load_t(&content, &mut reader, "x_embed.weight", &device)?;
239        let x_embed_b = load_t(&content, &mut reader, "x_embed.bias", &device)?;
240        let col_y_encoder_w = load_t(&content, &mut reader, "col_y_encoder.embedding.weight", &device)?;
241        let icl_y_encoder_w = load_t(&content, &mut reader, "icl_y_encoder.embedding.weight", &device)?;
242
243        let mut dist_embed = Vec::with_capacity(config.dist_embed_num_blocks);
244        for n in 0..config.dist_embed_num_blocks {
245            let p = format!("feature_distribution_embedder.layers.{n}");
246            let ind_vectors = load_t(&content, &mut reader, &format!("{p}.inducing_vectors"), &device)?;
247            let block1 = load_cross_attn_block(&content, &mut reader, &format!("{p}.cross_attn_block1"), &device, true)?;
248            let block2 = load_cross_attn_block(&content, &mut reader, &format!("{p}.cross_attn_block2"), &device, false)?;
249            dist_embed.push(IsabW { ind_vectors, block1, block2 });
250        }
251
252        let mut col_agg_blocks = Vec::with_capacity(config.feat_agg_num_blocks);
253        for n in 0..config.feat_agg_num_blocks {
254            let p = format!("column_aggregator.blocks.{n}");
255            col_agg_blocks.push(load_transformer_block(&content, &mut reader, &p, &device)?);
256        }
257        let col_agg_cls_tokens = load_t(&content, &mut reader, "column_aggregator.cls_tokens", &device)?;
258        let col_agg_out_ln_w = load_t(&content, &mut reader, "column_aggregator.out_ln.weight", &device)?;
259        let rope_freqs_t = load_t(&content, &mut reader, "column_aggregator.rope.freqs", &device)?;
260        let col_agg_rope_freqs: Vec<f32> = rope_freqs_t.flatten_all()?.to_vec1()?;
261
262        let mut icl_blocks = Vec::with_capacity(config.nlayers);
263        for n in 0..config.nlayers {
264            let p = format!("icl_blocks.{n}");
265            icl_blocks.push(load_icl_block(&content, &mut reader, &p, &device)?);
266        }
267
268        let output_norm_w = load_t(&content, &mut reader, "output_norm.weight", &device)?;
269
270        let decoder_ssmax_prefix = "many_class_decoder.softmax_scaling_layer";
271        let many_class_decoder = ManyClassDecoderW {
272            q_w: load_t(&content, &mut reader, "many_class_decoder.q_projection.weight", &device)?,
273            q_b: load_t(&content, &mut reader, "many_class_decoder.q_projection.bias", &device)?,
274            k_w: load_t(&content, &mut reader, "many_class_decoder.k_projection.weight", &device)?,
275            k_b: load_t(&content, &mut reader, "many_class_decoder.k_projection.bias", &device)?,
276            ssmax: if config.decoder_use_softmax_scaling {
277                Some(load_ssmax(&content, &mut reader, decoder_ssmax_prefix, &device)?)
278            } else {
279                None
280            },
281        };
282
283        Ok(Self {
284            device,
285            config,
286            x_embed_w,
287            x_embed_b,
288            col_y_encoder_w,
289            icl_y_encoder_w,
290            dist_embed,
291            col_agg_blocks,
292            col_agg_cls_tokens,
293            col_agg_rope_freqs,
294            col_agg_out_ln_w,
295            icl_blocks,
296            output_norm_w,
297            many_class_decoder,
298        })
299    }
300
301    // -----------------------------------------------------------------------
302    // Prediction
303    // -----------------------------------------------------------------------
304
305    /// Zero-shot classification. `y_support` are class indices (must satisfy
306    /// `max(y_support) < max_num_classes`). Returns probabilities `[n_query][n_classes]`.
307    pub fn predict_classification(
308        &self,
309        x_support: &[Vec<f32>],
310        y_support: &[usize],
311        x_query: &[Vec<f32>],
312        n_classes: usize,
313    ) -> Result<Vec<Vec<f32>>> {
314        let train_size = x_support.len();
315        let mut all_rows = x_support.to_vec();
316        all_rows.extend_from_slice(x_query);
317        let t = all_rows.len();
318        let h = all_rows[0].len();
319
320        let processed = preprocess_x(&all_rows, train_size);
321        let group_size = self.config.feature_group_size;
322        let grouped = feature_group_with_nan_indicators(&processed, h, group_size, self.config.use_nan_indicators);
323        let cell_dim = grouped[0][0].len();
324
325        let mut flat = vec![0f32; h * t * cell_dim];
326        for (row_idx, row) in grouped.iter().enumerate() {
327            for (col_idx, cell) in row.iter().enumerate() {
328                let base = col_idx * t * cell_dim + row_idx * cell_dim;
329                flat[base..base + cell_dim].copy_from_slice(cell);
330            }
331        }
332        let x_grouped = Tensor::from_vec(flat, (h, t, cell_dim), &self.device)?;
333
334        let y_col_emb = self.embedding_lookup(y_support, &self.col_y_encoder_w)?; // (train_size, embed_dim)
335        let col_out = self.dist_embed_forward(&x_grouped, &y_col_emb, train_size)?; // (H, T, embed_dim)
336        let row_out = self.col_agg_forward(&col_out, train_size)?; // (T, icl_dim)
337
338        let y_icl_emb = self.embedding_lookup(y_support, &self.icl_y_encoder_w)?; // (train_size, icl_dim)
339        let train_part = row_out.narrow(0, 0, train_size)?.broadcast_add(&y_icl_emb)?;
340        let mut r = if train_size < t {
341            Tensor::cat(&[&train_part, &row_out.narrow(0, train_size, t - train_size)?], 0)?
342        } else {
343            train_part
344        };
345        for blk in &self.icl_blocks {
346            r = self.icl_block_forward(&r, blk, train_size)?;
347        }
348        let r = zsfm_nn::rms_norm(&r, Some(&self.output_norm_w), RMS_EPS)?;
349
350        let train_emb = r.narrow(0, 0, train_size)?;
351        let test_emb = r.narrow(0, train_size, t - train_size)?;
352
353        let highest_target = *y_support.iter().max().context("y_support must be non-empty")?;
354        let logits = self.many_class_decoder_forward(&train_emb, &test_emb, y_support, highest_target, n_classes)?;
355
356        let flat_logits: Vec<f32> = logits.flatten_all()?.to_vec1()?;
357        let n_query = t - train_size;
358        const TEMPERATURE: f32 = 0.9;
359        let mut result = Vec::with_capacity(n_query);
360        for i in 0..n_query {
361            let row = &flat_logits[i * n_classes..(i + 1) * n_classes];
362            let scaled: Vec<f32> = row.iter().map(|&v| v / TEMPERATURE).collect();
363            result.push(softmax(&scaled));
364        }
365        Ok(result)
366    }
367
368    fn embedding_lookup(&self, y: &[usize], table: &Tensor) -> Result<Tensor> {
369        let dim = table.dim(1)?;
370        let idx = Tensor::from_vec(y.iter().map(|&c| c as u32).collect::<Vec<_>>(), y.len(), &self.device)?;
371        table.index_select(&idx, 0)?.reshape((y.len(), dim)).map_err(anyhow::Error::from)
372    }
373
374    /// `x_grouped`: `[H, T, cell_dim]`. `y_col_emb`: `[train_size, embed_dim]`. Returns `[H, T,
375    /// embed_dim]`.
376    fn dist_embed_forward(&self, x_grouped: &Tensor, y_col_emb: &Tensor, train_size: usize) -> Result<Tensor> {
377        let cell_emb = zsfm_nn::linear_bias(x_grouped, &self.x_embed_w, &self.x_embed_b)?;
378        let t = cell_emb.dim(1)?;
379        let train_part = cell_emb.narrow(1, 0, train_size)?.broadcast_add(&y_col_emb.unsqueeze(0)?)?;
380        let mut src = if train_size < t {
381            Tensor::cat(&[&train_part, &cell_emb.narrow(1, train_size, t - train_size)?], 1)?
382        } else {
383            train_part
384        };
385
386        let n_heads = self.config.dist_embed_num_heads;
387        let head_dim = self.config.embed_dim / n_heads;
388        for isab in &self.dist_embed {
389            let h = src.dim(0)?;
390            let num_inds = isab.ind_vectors.dim(0)?;
391            let dim = isab.ind_vectors.dim(1)?;
392            let ind = isab.ind_vectors.unsqueeze(0)?.expand((h, num_inds, dim))?.contiguous()?;
393            let kv_train = src.narrow(1, 0, train_size)?;
394            let hidden = cross_attn_block_forward(&ind, &kv_train, &isab.block1, n_heads, head_dim, train_size)?;
395            src = cross_attn_block_forward(&src, &hidden, &isab.block2, n_heads, head_dim, num_inds)?;
396        }
397        Ok(src)
398    }
399
400    /// `col_embeddings`: `[H, T, embed_dim]`. Returns row representations `[T, icl_dim]`.
401    fn col_agg_forward(&self, col_embeddings: &Tensor, _train_size: usize) -> Result<Tensor> {
402        let h = col_embeddings.dim(0)?;
403        let t = col_embeddings.dim(1)?;
404        let dim = self.config.embed_dim;
405        let n_cls = self.config.feat_agg_num_cls_tokens;
406        let n_heads = self.config.feat_agg_num_heads;
407        let head_dim = dim / n_heads;
408
409        let feat = col_embeddings.permute((1, 0, 2))?.contiguous()?; // (T, H, dim)
410        let cls = self.col_agg_cls_tokens.unsqueeze(0)?.expand((t, n_cls, dim))?.contiguous()?;
411        let mut seq = Tensor::cat(&[&cls, &feat], 1)?; // (T, H+C, dim)
412
413        let max_len = h + n_cls;
414        let (cos, sin) = rope_table(&self.col_agg_rope_freqs, max_len, &self.device)?;
415
416        let n_blocks = self.col_agg_blocks.len();
417        for (i, blk) in self.col_agg_blocks.iter().enumerate() {
418            if i + 1 == n_blocks {
419                let cls_q = seq.narrow(1, 0, n_cls)?;
420                seq = transformer_block_forward_cross(&cls_q, &seq, blk, n_heads, head_dim, Some((&cos, &sin)))?;
421            } else {
422                seq = transformer_block_forward(&seq, blk, n_heads, head_dim, Some((&cos, &sin)))?;
423            }
424        }
425
426        let out = zsfm_nn::rms_norm(&seq, Some(&self.col_agg_out_ln_w), RMS_EPS)?; // (T, C, dim)
427        out.reshape((t, n_cls * dim)).map_err(anyhow::Error::from)
428    }
429
430    /// One `ICLTransformerBlock`: train rows use full-head K/V, test rows use only the first
431    /// `icl_num_kv_heads_test` head(s) of K/V (broadcast to all query heads) — a GQA-style
432    /// reduction that applies to test rows *only*. `x`: `[T, icl_dim]`.
433    fn icl_block_forward(&self, x: &Tensor, blk: &IclBlockW, train_size: usize) -> Result<Tensor> {
434        let n_heads = self.config.icl_num_heads;
435        let icl_dim = self.config.icl_dim();
436        let head_dim = icl_dim / n_heads;
437        let t = x.dim(0)?;
438
439        let normed = zsfm_nn::rms_norm(x, Some(&blk.ln_w), RMS_EPS)?;
440        let q = zsfm_nn::linear_nobias(&normed, &blk.attn.q_w)?.reshape((t, n_heads, head_dim))?;
441        let x_train = normed.narrow(0, 0, train_size)?;
442        let k = zsfm_nn::linear_nobias(&x_train, &blk.attn.k_w)?.reshape((train_size, n_heads, head_dim))?;
443        let v = zsfm_nn::linear_nobias(&x_train, &blk.attn.v_w)?.reshape((train_size, n_heads, head_dim))?;
444
445        // (seq, heads, hd) -> (heads, seq, hd)
446        let q = q.permute((1, 0, 2))?.contiguous()?;
447        let k = k.permute((1, 0, 2))?.contiguous()?;
448        let v = v.permute((1, 0, 2))?.contiguous()?;
449
450        let attn_out = match self.config.icl_num_kv_heads_test {
451            Some(kv_heads_test) if train_size < t => {
452                let q_train = q.narrow(1, 0, train_size)?;
453                let q_test = q.narrow(1, train_size, t - train_size)?;
454                let out_train = sdpa_with_ssmax(&q_train, &k, &v, blk.attn.ssmax.as_ref(), train_size, n_heads, head_dim)?;
455
456                let k_test = k.narrow(0, 0, kv_heads_test)?;
457                let v_test = v.narrow(0, 0, kv_heads_test)?;
458                let k_test = repeat_heads(&k_test, n_heads / kv_heads_test)?;
459                let v_test = repeat_heads(&v_test, n_heads / kv_heads_test)?;
460                let out_test =
461                    sdpa_with_ssmax(&q_test, &k_test, &v_test, blk.attn.ssmax.as_ref(), train_size, n_heads, head_dim)?;
462                Tensor::cat(&[&out_train, &out_test], 1)? // (heads, T, hd)
463            }
464            _ => sdpa_with_ssmax(&q, &k, &v, blk.attn.ssmax.as_ref(), train_size, n_heads, head_dim)?,
465        };
466
467        let attn_out = attn_out.permute((1, 0, 2))?.contiguous()?.reshape((t, icl_dim))?;
468        let attn_out = zsfm_nn::linear_nobias(&attn_out, &blk.attn.out_w)?;
469        let x = (x + attn_out)?;
470
471        let ff_in = zsfm_nn::rms_norm(&x, Some(&blk.ln_mlp_w), RMS_EPS)?;
472        let ff = mlp_forward(&ff_in, &blk.mlp)?;
473        Ok((x + ff)?)
474    }
475
476    /// `train_emb`/`test_emb`: `[N, icl_dim]`/`[M, icl_dim]`. Returns logits `[M, n_classes]`.
477    fn many_class_decoder_forward(
478        &self,
479        train_emb: &Tensor,
480        test_emb: &Tensor,
481        y_support: &[usize],
482        highest_target: usize,
483        n_classes: usize,
484    ) -> Result<Tensor> {
485        let dec = &self.many_class_decoder;
486        let n_heads = self.config.decoder_num_heads;
487        let head_dim = self.config.decoder_head_dim;
488        let n = train_emb.dim(0)?;
489        let m = test_emb.dim(0)?;
490
491        let q = zsfm_nn::linear_bias(test_emb, &dec.q_w, &dec.q_b)?.reshape((m, n_heads, head_dim))?;
492        let k = zsfm_nn::linear_bias(train_emb, &dec.k_w, &dec.k_b)?.reshape((n, n_heads, head_dim))?;
493        let q = q.permute((1, 0, 2))?.contiguous()?; // (heads, M, hd)
494        let k = k.permute((1, 0, 2))?.contiguous()?; // (heads, N, hd)
495
496        let one_hot_width = highest_target + 1;
497        let mut one_hot = vec![0f32; n * one_hot_width];
498        for (i, &c) in y_support.iter().enumerate() {
499            one_hot[i * one_hot_width + c] = 1.0;
500        }
501        // V is shared identically across all heads (no v_projection in ManyClassDecoder).
502        let v_row = Tensor::from_vec(one_hot, (n, one_hot_width), &self.device)?;
503        let v = v_row.unsqueeze(0)?.expand((n_heads, n, one_hot_width))?.contiguous()?;
504
505        let num_chunks = one_hot_width.div_ceil(head_dim);
506        let padded_width = num_chunks * head_dim;
507        let v = if padded_width > one_hot_width {
508            v.pad_with_zeros(2, 0, padded_width - one_hot_width)?
509        } else {
510            v
511        };
512
513        let mut chunk_outs = Vec::with_capacity(num_chunks);
514        for c in 0..num_chunks {
515            let v_chunk = v.narrow(2, c * head_dim, head_dim)?;
516            let out = sdpa_with_ssmax(&q, &k, &v_chunk, dec.ssmax.as_ref(), n, n_heads, head_dim)?; // (heads, M, hd)
517            chunk_outs.push(out);
518        }
519        let out = if chunk_outs.len() == 1 {
520            chunk_outs.into_iter().next().unwrap()
521        } else {
522            let refs: Vec<&Tensor> = chunk_outs.iter().collect();
523            Tensor::cat(&refs, 2)? // (heads, M, num_chunks*hd)
524        };
525        let out = out.narrow(2, 0, one_hot_width)?; // (heads, M, one_hot_width)
526        let out = out.mean(0)?; // average over heads -> (M, one_hot_width)
527
528        let out = if n_classes > one_hot_width {
529            out.pad_with_zeros(1, 0, n_classes - one_hot_width)?
530        } else {
531            out.narrow(1, 0, n_classes)?
532        };
533
534        let clamped = out.clamp(1e-5f32, f32::INFINITY)?;
535        ((clamped + 3e-5)?).log().map_err(anyhow::Error::from)
536    }
537}
538
539/// Scaled dot-product attention with optional query-aware SSMax scaling. `q`/`k`/`v`:
540/// `[heads, seq, hd]`. `ssmax_n` is the KV sequence length used for SSMax's `log(n)` term.
541fn sdpa_with_ssmax(
542    q: &Tensor,
543    k: &Tensor,
544    v: &Tensor,
545    ssmax: Option<&SsmaxW>,
546    ssmax_n: usize,
547    n_heads: usize,
548    head_dim: usize,
549) -> Result<Tensor> {
550    let q = match ssmax {
551        Some(s) => apply_ssmax(q, s, ssmax_n, n_heads, head_dim)?,
552        None => q.clone(),
553    };
554    let scale = 1.0 / (head_dim as f64).sqrt();
555    let scores = (q.matmul(&k.transpose(1, 2)?)? * scale)?;
556    let probs = candle_nn::ops::softmax_last_dim(&scores)?;
557    probs.matmul(v).map_err(anyhow::Error::from)
558}
559
560/// `x`: `[kv_heads, seq, hd]` -> `[kv_heads * repeat, seq, hd]` (`repeat_interleave` along the
561/// head axis, matching torch's GQA broadcast on the math-backend path).
562fn repeat_heads(x: &Tensor, repeat: usize) -> Result<Tensor> {
563    if repeat == 1 {
564        return Ok(x.clone());
565    }
566    let (h, s, d) = x.dims3()?;
567    x.unsqueeze(1)?.expand((h, repeat, s, d))?.reshape((h * repeat, s, d)).map_err(anyhow::Error::from)
568}
569
570/// Query-aware elementwise softmax scaling (`SoftmaxScalingMLP`): `q * base_mlp(log n) * (1 +
571/// tanh(query_mlp(q)))`. `q`: `[heads, seq, hd]`.
572fn apply_ssmax(q: &Tensor, ssmax: &SsmaxW, n: usize, n_heads: usize, head_dim: usize) -> Result<Tensor> {
573    let device = q.device();
574    let logn = (n.max(1) as f32).ln();
575    let logn_t = Tensor::from_vec(vec![logn], (1, 1), device)?;
576    let base = zsfm_nn::linear_bias(&logn_t, &ssmax.base_fc1_w, &ssmax.base_fc1_b)?.gelu_erf()?;
577    let base = zsfm_nn::linear_bias(&base, &ssmax.base_fc2_w, &ssmax.base_fc2_b)?; // (1, n_heads*head_dim)
578    let base = base.reshape((n_heads, 1, head_dim))?;
579
580    // `q`'s leading dim is `batch * n_heads` flattened (batch = 1 for the ICL/decoder stages,
581    // batch = H columns for the distribution embedder's per-column ISAB) — `base` only varies
582    // per head, so it must be tiled across `batch` before broadcasting against `q`.
583    let batch_heads = q.dim(0)?;
584    let seq = q.dim(1)?;
585    let batch = batch_heads / n_heads;
586    let base = if batch > 1 {
587        base.unsqueeze(0)?.expand((batch, n_heads, 1, head_dim))?.reshape((batch_heads, 1, head_dim))?
588    } else {
589        base
590    };
591
592    let q_flat = q.reshape((batch_heads * seq, head_dim))?;
593    let qm = zsfm_nn::linear_bias(&q_flat, &ssmax.query_fc1_w, &ssmax.query_fc1_b)?.gelu_erf()?;
594    let qm = zsfm_nn::linear_bias(&qm, &ssmax.query_fc2_w, &ssmax.query_fc2_b)?;
595    let modulation = (qm.tanh()? + 1.0)?.reshape((batch_heads, seq, head_dim))?;
596
597    let scales = base.broadcast_mul(&modulation)?;
598    Ok(q.broadcast_mul(&scales)?)
599}
600
601fn mlp_forward(x: &Tensor, mlp: &MlpW) -> Result<Tensor> {
602    let h = zsfm_nn::linear_nobias(x, &mlp.fc1_w)?.gelu_erf()?;
603    zsfm_nn::linear_nobias(&h, &mlp.fc2_w).map_err(anyhow::Error::from)
604}
605
606/// `CrossAttentionBlock`: `x = query + Attn(LN_q(query), LN_kv(context))`, then `x = x +
607/// MLP(LN2(x))`. `query`/`context`: `[batch, seq, dim]`, `batch` is the induced-self-attention
608/// block's "H columns as batch" axis.
609fn cross_attn_block_forward(
610    query: &Tensor,
611    context: &Tensor,
612    blk: &CrossAttnBlockW,
613    n_heads: usize,
614    head_dim: usize,
615    ssmax_n: usize,
616) -> Result<Tensor> {
617    let q_normed = zsfm_nn::rms_norm(query, Some(&blk.ln_q_w), RMS_EPS)?;
618    let kv_normed = zsfm_nn::rms_norm(context, Some(&blk.ln_kv_w), RMS_EPS)?;
619    let attn_out = batched_qkv_attention(&q_normed, &kv_normed, &blk.attn, n_heads, head_dim, ssmax_n, None)?;
620    let x = (query + attn_out)?;
621    let ff_in = zsfm_nn::rms_norm(&x, Some(&blk.ln2_w), RMS_EPS)?;
622    let ff = mlp_forward(&ff_in, &blk.mlp)?;
623    Ok((x + ff)?)
624}
625
626/// `TransformerBlock` self-attention: `x = x + Attn(LN(x))`, then `x = x + MLP(LN_mlp(x))`.
627/// `x`: `[batch, seq, dim]`, `batch` is the "T rows as batch" axis (`ColumnAggregator`).
628fn transformer_block_forward(
629    x: &Tensor,
630    blk: &TransformerBlockW,
631    n_heads: usize,
632    head_dim: usize,
633    rope_cos_sin: Option<(&Tensor, &Tensor)>,
634) -> Result<Tensor> {
635    let normed = zsfm_nn::rms_norm(x, Some(&blk.ln_w), RMS_EPS)?;
636    let attn_out = batched_qkv_self_attention(&normed, &blk.attn, n_heads, head_dim, rope_cos_sin)?;
637    let x = (x + attn_out)?;
638    let ff_in = zsfm_nn::rms_norm(&x, Some(&blk.ln_mlp_w), RMS_EPS)?;
639    let ff = mlp_forward(&ff_in, &blk.mlp)?;
640    Ok((x + ff)?)
641}
642
643/// `TransformerBlock.forward_cross`: CLS-tokens-as-query readout. The *same* `layernorm` weight
644/// normalizes both the query and the context (matching the reference's literal
645/// `self.layernorm(...)` reuse on both sides — not a separate `layernorm_kv`).
646fn transformer_block_forward_cross(
647    query: &Tensor,
648    context: &Tensor,
649    blk: &TransformerBlockW,
650    n_heads: usize,
651    head_dim: usize,
652    rope_cos_sin: Option<(&Tensor, &Tensor)>,
653) -> Result<Tensor> {
654    let q_normed = zsfm_nn::rms_norm(query, Some(&blk.ln_w), RMS_EPS)?;
655    let kv_normed = zsfm_nn::rms_norm(context, Some(&blk.ln_w), RMS_EPS)?;
656    let attn_out = batched_qkv_attention(&q_normed, &kv_normed, &blk.attn, n_heads, head_dim, 0, rope_cos_sin)?;
657    let x = (query + attn_out)?;
658    let ff_in = zsfm_nn::rms_norm(&x, Some(&blk.ln_mlp_w), RMS_EPS)?;
659    let ff = mlp_forward(&ff_in, &blk.mlp)?;
660    Ok((x + ff)?)
661}
662
663/// Cross-attention over an explicit `[batch, seq, dim]` batch axis (the induced-self-attention
664/// "H columns as batch" axis, or `ColumnAggregator`'s "T rows as batch" axis for the CLS
665/// readout, where RoPE applies to both Q and K — `forward_cross` rotates both sides using each
666/// tensor's own sequence length, so the CLS query and the full K sequence get their true
667/// absolute positions).
668fn batched_qkv_attention(
669    q_in: &Tensor,
670    kv_in: &Tensor,
671    attn: &QkvW,
672    n_heads: usize,
673    head_dim: usize,
674    ssmax_n: usize,
675    rope_cos_sin: Option<(&Tensor, &Tensor)>,
676) -> Result<Tensor> {
677    let dim = n_heads * head_dim;
678    let batch = q_in.dim(0)?;
679    let q_len = q_in.dim(1)?;
680    let kv_len = kv_in.dim(1)?;
681
682    let q = zsfm_nn::linear_nobias(q_in, &attn.q_w)?.reshape((batch, q_len, n_heads, head_dim))?;
683    let k = zsfm_nn::linear_nobias(kv_in, &attn.k_w)?.reshape((batch, kv_len, n_heads, head_dim))?;
684    let v = zsfm_nn::linear_nobias(kv_in, &attn.v_w)?.reshape((batch, kv_len, n_heads, head_dim))?;
685
686    let q = q.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch * n_heads, q_len, head_dim))?;
687    let k = k.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch * n_heads, kv_len, head_dim))?;
688    let v = v.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch * n_heads, kv_len, head_dim))?;
689
690    let (q, k) = if let Some((cos, sin)) = rope_cos_sin {
691        (apply_rope(&q, cos, sin, batch, n_heads)?, apply_rope(&k, cos, sin, batch, n_heads)?)
692    } else {
693        (q, k)
694    };
695
696    let n = if ssmax_n > 0 { ssmax_n } else { kv_len };
697    let out = sdpa_with_ssmax(&q, &k, &v, attn.ssmax.as_ref(), n, n_heads, head_dim)?;
698    let out = out.reshape((batch, n_heads, q_len, head_dim))?.permute((0, 2, 1, 3))?.contiguous()?.reshape((
699        batch,
700        q_len,
701        dim,
702    ))?;
703    zsfm_nn::linear_nobias(&out, &attn.out_w).map_err(anyhow::Error::from)
704}
705
706/// Self-attention with optional RoPE over an explicit `[batch, seq, dim]` batch axis.
707fn batched_qkv_self_attention(
708    x: &Tensor,
709    attn: &QkvW,
710    n_heads: usize,
711    head_dim: usize,
712    rope_cos_sin: Option<(&Tensor, &Tensor)>,
713) -> Result<Tensor> {
714    let dim = n_heads * head_dim;
715    let batch = x.dim(0)?;
716    let seq = x.dim(1)?;
717
718    let q = zsfm_nn::linear_nobias(x, &attn.q_w)?.reshape((batch, seq, n_heads, head_dim))?;
719    let k = zsfm_nn::linear_nobias(x, &attn.k_w)?.reshape((batch, seq, n_heads, head_dim))?;
720    let v = zsfm_nn::linear_nobias(x, &attn.v_w)?.reshape((batch, seq, n_heads, head_dim))?;
721
722    let q = q.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch * n_heads, seq, head_dim))?;
723    let k = k.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch * n_heads, seq, head_dim))?;
724    let v = v.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch * n_heads, seq, head_dim))?;
725
726    let (q, k) = if let Some((cos, sin)) = rope_cos_sin {
727        (apply_rope(&q, cos, sin, batch, n_heads)?, apply_rope(&k, cos, sin, batch, n_heads)?)
728    } else {
729        (q, k)
730    };
731
732    let out = sdpa_with_ssmax(&q, &k, &v, attn.ssmax.as_ref(), seq, n_heads, head_dim)?;
733    let out =
734        out.reshape((batch, n_heads, seq, head_dim))?.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch, seq, dim))?;
735    zsfm_nn::linear_nobias(&out, &attn.out_w).map_err(anyhow::Error::from)
736}
737
738/// Non-interleaved RoPE. `x`: `[batch*heads, seq, head_dim]`. `cos`/`sin`: `[max_len, head_dim]`.
739fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor, _batch: usize, _n_heads: usize) -> Result<Tensor> {
740    let seq = x.dim(1)?;
741    let head_dim = x.dim(2)?;
742    let half = head_dim / 2;
743    let cos = cos.narrow(0, 0, seq)?.unsqueeze(0)?; // (1, seq, hd)
744    let sin = sin.narrow(0, 0, seq)?.unsqueeze(0)?;
745    let x1 = x.narrow(2, 0, half)?;
746    let x2 = x.narrow(2, half, half)?;
747    let rotated = Tensor::cat(&[&x2.neg()?, &x1], 2)?;
748    Ok((x.broadcast_mul(&cos)? + rotated.broadcast_mul(&sin)?)?)
749}
750
751fn rope_table(freqs: &[f32], max_len: usize, device: &Device) -> Result<(Tensor, Tensor)> {
752    let half = freqs.len();
753    let mut cos = vec![0f32; max_len * half * 2];
754    let mut sin = vec![0f32; max_len * half * 2];
755    for p in 0..max_len {
756        for i in 0..half {
757            let angle = p as f32 * freqs[i];
758            let (s, c) = angle.sin_cos();
759            cos[p * 2 * half + i] = c;
760            cos[p * 2 * half + half + i] = c;
761            sin[p * 2 * half + i] = s;
762            sin[p * 2 * half + half + i] = s;
763        }
764    }
765    let cos_t = Tensor::from_vec(cos, (max_len, 2 * half), device)?;
766    let sin_t = Tensor::from_vec(sin, (max_len, 2 * half), device)?;
767    Ok((cos_t, sin_t))
768}
769
770fn softmax(logits: &[f32]) -> Vec<f32> {
771    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
772    let exps: Vec<f32> = logits.iter().map(|&v| (v - max).exp()).collect();
773    let sum: f32 = exps.iter().sum();
774    exps.into_iter().map(|v| v / sum).collect()
775}
776
777// ---------------------------------------------------------------------------
778// Preprocessing
779// ---------------------------------------------------------------------------
780
781/// Mean-impute (no-op here — this port assumes no NaNs, matching the scope decision established
782/// for every other model in this workspace) + standardize with `ddof=1` (sample std, matching
783/// `torch_nanstd`'s `N-1` correction — *not* the `ddof=0` used by Mitra/TabDPT/TabICL's simpler
784/// preprocessing), `+ eps` before dividing, then clip to `[-100, 100]`.
785fn preprocess_x(rows: &[Vec<f32>], train_size: usize) -> Vec<Vec<f32>> {
786    let n_feat = rows[0].len();
787    let mut mean = vec![0f32; n_feat];
788    let mut std = vec![0f32; n_feat];
789    for f in 0..n_feat {
790        let m: f32 = rows[..train_size].iter().map(|r| r[f]).sum::<f32>() / train_size as f32;
791        let denom = if train_size > 1 { (train_size - 1) as f32 } else { 1.0 };
792        let var: f32 = rows[..train_size].iter().map(|r| (r[f] - m).powi(2)).sum::<f32>() / denom;
793        mean[f] = m;
794        std[f] = if var == 0.0 || train_size <= 1 { 1.0 } else { var.sqrt() };
795    }
796    let eps = f32::EPSILON;
797    rows.iter()
798        .map(|row| {
799            row.iter().enumerate().map(|(f, &v)| ((v - mean[f]) / (std[f] + eps)).clamp(-100.0, 100.0)).collect()
800        })
801        .collect()
802}
803
804/// Circular-permutation feature grouping (same formula as TabICL's `feature_group_same`: group
805/// `g`'s values are input features at offsets `2^0, 2^1, .., 2^(size-1)` past `g`, mod `H`),
806/// with NaN/Inf indicator features (always `0.0` in this no-NaN-support port) concatenated after
807/// the real grouped values, matching `x_grouped = cat([x_grouped, ind_grouped], dim=-1)`.
808fn feature_group_with_nan_indicators(rows: &[Vec<f32>], h: usize, size: usize, use_nan_indicators: bool) -> Vec<Vec<Vec<f32>>> {
809    rows.iter()
810        .map(|row| {
811            (0..h)
812                .map(|g| {
813                    let mut cell: Vec<f32> = (0..size).map(|k| row[(g + (1usize << k)) % h]).collect();
814                    if use_nan_indicators {
815                        cell.extend(std::iter::repeat_n(0.0f32, size));
816                    }
817                    cell
818                })
819                .collect()
820        })
821        .collect()
822}
zsfm_tabpfn/infer/mod.rs

zsfm_tabpfn/infer/
mod.rs
