zsfm_mitra/infer/
mod.rs

1//! Mitra (Tab2D) inference engine — zero-shot forward pass only (no fine-tuning; see the crate
2//! README for why). Architecture: per-feature quantile-bucketize embedding → prepend a learned
3//! y-embedding as an extra "feature" column → 12 layers of (row self/cross-attention → MLP →
4//! feature self-attention → MLP) → final LayerNorm + linear head, read out at the y-column.
5//!
6//! Two deliberate, documented divergences from AutoGluon's default `MitraClassifier`/
7//! `MitraRegressor` (both scope decisions, not omissions):
8//! - `random_mirror_x`/`random_mirror_regression` (default ON upstream) are OFF here for
9//!   determinism — even upstream they draw from the *unseeded* global NumPy RNG, so upstream's
10//!   own "default" behavior isn't reproducible either without external global-seed control.
11//! - The per-call support-row shuffle (`np.random.RandomState.choice`, seeded) is not
12//!   replicated — same scope call already made for TabFM's OOF K-fold splitting (see
13//!   `tabfm/src/ensemble/oof.rs`): porting NumPy's legacy `RandomState` exactly is extra work
14//!   for a step the model is mathematically invariant to (attention has no row positional
15//!   encoding); rows are fed in the given order.
16
17use std::io::{BufReader, Read, Seek};
18use std::path::Path;
19
20use anyhow::{Context, Result};
21use candle_core::quantized::gguf_file;
22use candle_core::{DType, Device, Tensor};
23
24use crate::config::{MitraConfig, Task};
25
26// ---------------------------------------------------------------------------
27// Weight structs
28// ---------------------------------------------------------------------------
29
30struct AttnW {
31    q_w: Tensor,
32    q_b: Tensor,
33    k_w: Tensor,
34    k_b: Tensor,
35    v_w: Tensor,
36    v_b: Tensor,
37    o_w: Tensor,
38    o_b: Tensor,
39}
40
41struct BlockW {
42    ln1_w: Tensor,
43    ln1_b: Tensor,
44    ln2_w: Tensor,
45    ln2_b: Tensor,
46    ln3_w: Tensor,
47    ln3_b: Tensor,
48    ln4_w: Tensor,
49    ln4_b: Tensor,
50    attn_row: AttnW,
51    attn_feat: AttnW,
52    mlp1_fc1_w: Tensor,
53    mlp1_fc1_b: Tensor,
54    mlp1_fc2_w: Tensor,
55    mlp1_fc2_b: Tensor,
56    mlp2_fc1_w: Tensor,
57    mlp2_fc1_b: Tensor,
58    mlp2_fc2_w: Tensor,
59    mlp2_fc2_b: Tensor,
60}
61
62pub struct MitraModel {
63    device: Device,
64    config: MitraConfig,
65    x_embed_w: Tensor, // [dim, 1]
66    x_embed_b: Tensor, // [dim]
67    y_embed_w: Tensor, // classifier: [dim_output, dim] (Embedding table); regressor: [dim, 1] (Linear weight)
68    y_embed_b: Option<Tensor>, // regressor only: [dim]
69    y_mask_w: Tensor,  // [1, dim]
70    blocks: Vec<BlockW>,
71    norm_f_w: Tensor,
72    norm_f_b: Tensor,
73    head_w: Tensor,
74    head_b: Tensor,
75}
76
77const LN_EPS: f64 = 1e-5;
78
79// ---------------------------------------------------------------------------
80// GGUF loading
81// ---------------------------------------------------------------------------
82
83fn load_t(content: &gguf_file::Content, reader: &mut (impl Read + Seek), name: &str, device: &Device) -> Result<Tensor> {
84    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
85}
86
87fn load_attn(
88    content: &gguf_file::Content,
89    reader: &mut (impl Read + Seek),
90    prefix: &str,
91    device: &Device,
92) -> Result<AttnW> {
93    let mut t = |s: &str| load_t(content, reader, &format!("{prefix}.{s}"), device);
94    Ok(AttnW {
95        q_w: t("q.weight")?,
96        q_b: t("q.bias")?,
97        k_w: t("k.weight")?,
98        k_b: t("k.bias")?,
99        v_w: t("v.weight")?,
100        v_b: t("v.bias")?,
101        o_w: t("o.weight")?,
102        o_b: t("o.bias")?,
103    })
104}
105
106impl MitraModel {
107    pub fn load(gguf_path: &Path, config: MitraConfig) -> Result<Self> {
108        let device = Device::Cpu;
109        let file = std::fs::File::open(gguf_path).with_context(|| format!("open {}", gguf_path.display()))?;
110        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
111        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
112
113        let x_embed_w = load_t(&content, &mut reader, "x_embed.weight", &device)?;
114        let x_embed_b = load_t(&content, &mut reader, "x_embed.bias", &device)?;
115        let y_embed_w = load_t(&content, &mut reader, "y_embed.weight", &device)?;
116        let y_embed_b = match config.task {
117            Task::Regression => Some(load_t(&content, &mut reader, "y_embed.bias", &device)?),
118            Task::Classification => None,
119        };
120        let y_mask_w = load_t(&content, &mut reader, "y_mask.weight", &device)?;
121
122        let mut blocks = Vec::with_capacity(config.n_layers);
123        for n in 0..config.n_layers {
124            let p = |s: &str| format!("blk.{n}.{s}");
125            let attn_row = load_attn(&content, &mut reader, &p("attn_row"), &device)?;
126            let attn_feat = load_attn(&content, &mut reader, &p("attn_feat"), &device)?;
127            blocks.push(BlockW {
128                ln1_w: load_t(&content, &mut reader, &p("ln1.weight"), &device)?,
129                ln1_b: load_t(&content, &mut reader, &p("ln1.bias"), &device)?,
130                ln2_w: load_t(&content, &mut reader, &p("ln2.weight"), &device)?,
131                ln2_b: load_t(&content, &mut reader, &p("ln2.bias"), &device)?,
132                ln3_w: load_t(&content, &mut reader, &p("ln3.weight"), &device)?,
133                ln3_b: load_t(&content, &mut reader, &p("ln3.bias"), &device)?,
134                ln4_w: load_t(&content, &mut reader, &p("ln4.weight"), &device)?,
135                ln4_b: load_t(&content, &mut reader, &p("ln4.bias"), &device)?,
136                attn_row,
137                attn_feat,
138                mlp1_fc1_w: load_t(&content, &mut reader, &p("mlp1_fc1.weight"), &device)?,
139                mlp1_fc1_b: load_t(&content, &mut reader, &p("mlp1_fc1.bias"), &device)?,
140                mlp1_fc2_w: load_t(&content, &mut reader, &p("mlp1_fc2.weight"), &device)?,
141                mlp1_fc2_b: load_t(&content, &mut reader, &p("mlp1_fc2.bias"), &device)?,
142                mlp2_fc1_w: load_t(&content, &mut reader, &p("mlp2_fc1.weight"), &device)?,
143                mlp2_fc1_b: load_t(&content, &mut reader, &p("mlp2_fc1.bias"), &device)?,
144                mlp2_fc2_w: load_t(&content, &mut reader, &p("mlp2_fc2.weight"), &device)?,
145                mlp2_fc2_b: load_t(&content, &mut reader, &p("mlp2_fc2.bias"), &device)?,
146            });
147        }
148
149        let norm_f_w = load_t(&content, &mut reader, "norm_f.weight", &device)?;
150        let norm_f_b = load_t(&content, &mut reader, "norm_f.bias", &device)?;
151        let head_w = load_t(&content, &mut reader, "head.weight", &device)?;
152        let head_b = load_t(&content, &mut reader, "head.bias", &device)?;
153
154        Ok(Self {
155            device,
156            config,
157            x_embed_w,
158            x_embed_b,
159            y_embed_w,
160            y_embed_b,
161            y_mask_w,
162            blocks,
163            norm_f_w,
164            norm_f_b,
165            head_w,
166            head_b,
167        })
168    }
169
170    // -----------------------------------------------------------------------
171    // Prediction
172    // -----------------------------------------------------------------------
173
174    /// Zero-shot classification. `y_support` are class indices `0..n_classes`. Returns raw
175    /// logits `[n_query][n_classes]` (softmax if you want probabilities).
176    pub fn predict_classification(
177        &self,
178        x_support: &[Vec<f32>],
179        y_support: &[usize],
180        x_query: &[Vec<f32>],
181        n_classes: usize,
182    ) -> Result<Vec<Vec<f32>>> {
183        anyhow::ensure!(self.config.task == Task::Classification, "model was loaded as a regressor");
184        let (x_s, x_q, kept) = preprocess_x(x_support, x_query);
185        let n_s = x_s.len();
186        let n_q = x_q.len();
187        let n_feat = kept.len();
188
189        let x_emb = self.embed_x(&x_s, &x_q, n_s, n_q, n_feat)?;
190        let y_support_emb = self.embed_y_classes_support(y_support, n_s)?;
191        let y_query_mask = self.embed_y_mask(n_q)?;
192
193        let logits = self.forward(x_emb, y_support_emb, y_query_mask, n_s, n_q, n_feat)?;
194        let logits: Vec<f32> = logits.flatten_all()?.to_vec1()?;
195        let dim_out = self.config.dim_output;
196        Ok((0..n_q).map(|i| logits[i * dim_out..i * dim_out + n_classes].to_vec()).collect())
197    }
198
199    /// Zero-shot regression. Returns predicted values in `y_support`'s original scale.
200    pub fn predict_regression(
201        &self,
202        x_support: &[Vec<f32>],
203        y_support: &[f32],
204        x_query: &[Vec<f32>],
205    ) -> Result<Vec<f32>> {
206        anyhow::ensure!(self.config.task == Task::Regression, "model was loaded as a classifier");
207        let (x_s, x_q, kept) = preprocess_x(x_support, x_query);
208        let n_s = x_s.len();
209        let n_q = x_q.len();
210        let n_feat = kept.len();
211
212        let y_min = y_support.iter().copied().fold(f32::INFINITY, f32::min);
213        let y_max = y_support.iter().copied().fold(f32::NEG_INFINITY, f32::max);
214        anyhow::ensure!(y_max > y_min, "y_support must have at least two distinct values");
215        let y_scaled: Vec<f32> = y_support.iter().map(|&v| (v - y_min) / (y_max - y_min)).collect();
216
217        let x_emb = self.embed_x(&x_s, &x_q, n_s, n_q, n_feat)?;
218        let y_support_emb = self.embed_y_regression_support(&y_scaled)?;
219        let y_query_mask = self.embed_y_mask(n_q)?;
220
221        let out = self.forward(x_emb, y_support_emb, y_query_mask, n_s, n_q, n_feat)?;
222        let out: Vec<f32> = out.flatten_all()?.to_vec1()?;
223        Ok(out.into_iter().map(|v| v * (y_max - y_min) + y_min).collect())
224    }
225
226    fn embed_x(&self, x_s: &[Vec<f32>], x_q: &[Vec<f32>], n_s: usize, n_q: usize, n_feat: usize) -> Result<(Tensor, Tensor)> {
227        let (bx_s, bx_q) = quantile_bucketize_normalize(x_s, x_q);
228        let flat_s: Vec<f32> = bx_s.into_iter().flatten().collect();
229        let flat_q: Vec<f32> = bx_q.into_iter().flatten().collect();
230
231        let x_s_t = Tensor::from_vec(flat_s, (n_s * n_feat, 1), &self.device)?;
232        let x_q_t = Tensor::from_vec(flat_q, (n_q * n_feat, 1), &self.device)?;
233
234        let x_s_emb = zsfm_nn::linear_bias(&x_s_t, &self.x_embed_w, &self.x_embed_b)?.reshape((n_s, n_feat, self.config.dim))?;
235        let x_q_emb = zsfm_nn::linear_bias(&x_q_t, &self.x_embed_w, &self.x_embed_b)?.reshape((n_q, n_feat, self.config.dim))?;
236        Ok((x_s_emb, x_q_emb))
237    }
238
239    fn embed_y_classes_support(&self, y_support: &[usize], n_s: usize) -> Result<Tensor> {
240        let dim = self.config.dim;
241        let table: Vec<f32> = self.y_embed_w.flatten_all()?.to_vec1()?;
242        let mut out = vec![0f32; n_s * dim];
243        for (i, &cls) in y_support.iter().enumerate() {
244            out[i * dim..(i + 1) * dim].copy_from_slice(&table[cls * dim..(cls + 1) * dim]);
245        }
246        Ok(Tensor::from_vec(out, (n_s, 1, dim), &self.device)?)
247    }
248
249    fn embed_y_regression_support(&self, y_scaled: &[f32]) -> Result<Tensor> {
250        let n_s = y_scaled.len();
251        let y_t = Tensor::from_vec(y_scaled.to_vec(), (n_s, 1), &self.device)?;
252        let b = self.y_embed_b.as_ref().context("regressor y_embed missing bias")?;
253        let emb = zsfm_nn::linear_bias(&y_t, &self.y_embed_w, b)?; // (n_s, dim)
254        Ok(emb.reshape((n_s, 1, self.config.dim))?)
255    }
256
257    fn embed_y_mask(&self, n_q: usize) -> Result<Tensor> {
258        let dim = self.config.dim;
259        let row: Vec<f32> = self.y_mask_w.flatten_all()?.to_vec1()?; // [dim]
260        let mut out = Vec::with_capacity(n_q * dim);
261        for _ in 0..n_q {
262            out.extend_from_slice(&row);
263        }
264        Ok(Tensor::from_vec(out, (n_q, 1, dim), &self.device)?)
265    }
266
267    /// `x_emb` = (support_x_emb, query_x_emb), each `[n, n_feat, dim]`.
268    /// `y_support_emb`: `[n_s, 1, dim]`. `y_query_mask`: `[n_q, 1, dim]`.
269    /// Returns `[n_q, dim_output]` (the model's output read out at the y-column).
270    fn forward(
271        &self,
272        x_emb: (Tensor, Tensor),
273        y_support_emb: Tensor,
274        y_query_mask: Tensor,
275        n_s: usize,
276        n_q: usize,
277        n_feat: usize,
278    ) -> Result<Tensor> {
279        let (x_s_emb, x_q_emb) = x_emb;
280        // Pack: y at column 0, features at columns 1..=n_feat.
281        let mut support = Tensor::cat(&[&y_support_emb, &x_s_emb], 1)?; // (n_s, f+1, dim)
282        let mut query = Tensor::cat(&[&y_query_mask, &x_q_emb], 1)?; // (n_q, f+1, dim)
283        let f1 = n_feat + 1;
284
285        for blk in &self.blocks {
286            let (s2, q2) = query_block_forward(&support, &query, blk, self.config.n_heads, n_s, n_q, f1)?;
287            support = s2;
288            query = q2;
289        }
290
291        let query = zsfm_nn::layer_norm(&query, &self.norm_f_w, &self.norm_f_b, LN_EPS)?;
292        let query = zsfm_nn::linear_bias(&query, &self.head_w, &self.head_b)?; // (n_q, f+1, dim_output)
293        // Column 0 is the y-slot.
294        query.narrow(1, 0, 1)?.reshape((n_q, self.config.dim_output))
295            .map_err(anyhow::Error::from)
296    }
297}
298
299/// One full `Layer` (row-attn → MLP → feature-attn → MLP) applied to both `support` and
300/// `query` streams, mirroring `Tab2D`'s Python `Layer.forward` CPU path exactly.
301fn query_block_forward(
302    support: &Tensor,
303    query: &Tensor,
304    blk: &BlockW,
305    n_heads: usize,
306    n_s: usize,
307    n_q: usize,
308    f1: usize,
309) -> Result<(Tensor, Tensor)> {
310    // --- Row attention (across observations; query cross-attends to support) ---
311    let res_s = support.clone();
312    let res_q = query.clone();
313    let s_ln = zsfm_nn::layer_norm(support, &blk.ln1_w, &blk.ln1_b, LN_EPS)?;
314    let q_ln = zsfm_nn::layer_norm(query, &blk.ln1_w, &blk.ln1_b, LN_EPS)?;
315
316    // (n, f, d) -> (f, n, d): attention batch = feature columns, sequence = observations.
317    let s_row = s_ln.permute((1, 0, 2))?.contiguous()?;
318    let q_row = q_ln.permute((1, 0, 2))?.contiguous()?;
319
320    let s_att = mha(&s_row, &s_row, &s_row, &blk.attn_row, n_heads)?;
321    let q_att = mha(&q_row, &s_row, &s_row, &blk.attn_row, n_heads)?;
322
323    let s_att = s_att.permute((1, 0, 2))?.contiguous()?.reshape((n_s, f1, s_att.dim(2)?))?;
324    let q_att = q_att.permute((1, 0, 2))?.contiguous()?.reshape((n_q, f1, q_att.dim(2)?))?;
325
326    let mut support = (res_s + s_att)?;
327    let mut query = (res_q + q_att)?;
328
329    // --- MLP 1 ---
330    let res_s = support.clone();
331    let res_q = query.clone();
332    let s_ln = zsfm_nn::layer_norm(&support, &blk.ln2_w, &blk.ln2_b, LN_EPS)?;
333    let q_ln = zsfm_nn::layer_norm(&query, &blk.ln2_w, &blk.ln2_b, LN_EPS)?;
334    let s_mlp = zsfm_nn::linear_bias(&zsfm_nn::linear_bias(&s_ln, &blk.mlp1_fc1_w, &blk.mlp1_fc1_b)?.gelu_erf()?, &blk.mlp1_fc2_w, &blk.mlp1_fc2_b)?;
335    let q_mlp = zsfm_nn::linear_bias(&zsfm_nn::linear_bias(&q_ln, &blk.mlp1_fc1_w, &blk.mlp1_fc1_b)?.gelu_erf()?, &blk.mlp1_fc2_w, &blk.mlp1_fc2_b)?;
336    support = (res_s + s_mlp)?;
337    query = (res_q + q_mlp)?;
338
339    // --- Feature attention (across columns, per observation; no cross term) ---
340    let res_s = support.clone();
341    let res_q = query.clone();
342    let s_ln = zsfm_nn::layer_norm(&support, &blk.ln3_w, &blk.ln3_b, LN_EPS)?; // already (n_s, f1, d)
343    let q_ln = zsfm_nn::layer_norm(&query, &blk.ln3_w, &blk.ln3_b, LN_EPS)?; // already (n_q, f1, d)
344
345    let s_att = mha(&s_ln, &s_ln, &s_ln, &blk.attn_feat, n_heads)?;
346    let q_att = mha(&q_ln, &q_ln, &q_ln, &blk.attn_feat, n_heads)?;
347
348    support = (res_s + s_att)?;
349    query = (res_q + q_att)?;
350
351    // --- MLP 2 ---
352    let res_s = support.clone();
353    let res_q = query.clone();
354    let s_ln = zsfm_nn::layer_norm(&support, &blk.ln4_w, &blk.ln4_b, LN_EPS)?;
355    let q_ln = zsfm_nn::layer_norm(&query, &blk.ln4_w, &blk.ln4_b, LN_EPS)?;
356    let s_mlp = zsfm_nn::linear_bias(&zsfm_nn::linear_bias(&s_ln, &blk.mlp2_fc1_w, &blk.mlp2_fc1_b)?.gelu_erf()?, &blk.mlp2_fc2_w, &blk.mlp2_fc2_b)?;
357    let q_mlp = zsfm_nn::linear_bias(&zsfm_nn::linear_bias(&q_ln, &blk.mlp2_fc1_w, &blk.mlp2_fc1_b)?.gelu_erf()?, &blk.mlp2_fc2_w, &blk.mlp2_fc2_b)?;
358    support = (res_s + s_mlp)?;
359    query = (res_q + q_mlp)?;
360
361    Ok((support, query))
362}
363
364/// Standard scaled-dot-product multi-head attention. `q`: `[batch, sq, dim]`, `k`/`v`:
365/// `[batch, skv, dim]`. Uses the default PyTorch `scaled_dot_product_attention` scale of
366/// `1/sqrt(head_dim)`.
367fn mha(q: &Tensor, k: &Tensor, v: &Tensor, w: &AttnW, n_heads: usize) -> Result<Tensor> {
368    let (batch, sq, dim) = q.dims3()?;
369    let skv = k.dim(1)?;
370    let head_dim = dim / n_heads;
371
372    let q = zsfm_nn::linear_bias(q, &w.q_w, &w.q_b)?;
373    let k = zsfm_nn::linear_bias(k, &w.k_w, &w.k_b)?;
374    let v = zsfm_nn::linear_bias(v, &w.v_w, &w.v_b)?;
375
376    let q = q.reshape((batch, sq, n_heads, head_dim))?.permute((0, 2, 1, 3))?.contiguous()?;
377    let k = k.reshape((batch, skv, n_heads, head_dim))?.permute((0, 2, 1, 3))?.contiguous()?;
378    let v = v.reshape((batch, skv, n_heads, head_dim))?.permute((0, 2, 1, 3))?.contiguous()?;
379
380    let scale = (head_dim as f64).sqrt();
381    let scores = (q.matmul(&k.transpose(2, 3)?)? / scale)?;
382    let attn = candle_nn::ops::softmax_last_dim(&scores)?;
383    let out = attn.matmul(&v)?; // (batch, h, sq, head_dim)
384    let out = out.permute((0, 2, 1, 3))?.contiguous()?.reshape((batch, sq, dim))?;
385    zsfm_nn::linear_bias(&out, &w.o_w, &w.o_b)
386}
387
388// ---------------------------------------------------------------------------
389// Preprocessing (Preprocessor.fit/transform_X, minus flags that default off:
390// use_quantile_transformer, use_feature_count_scaling, use_random_transforms,
391// shuffle_features, random_mirror_x)
392// ---------------------------------------------------------------------------
393
394/// Mean-impute NaNs (mean computed pre-imputation from `x_support`), then drop feature columns
395/// that are constant across `x_support` (both computed from `x_support`, applied to both).
396/// Returns `(x_support, x_query, kept_feature_indices)`.
397fn preprocess_x(x_support: &[Vec<f32>], x_query: &[Vec<f32>]) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<usize>) {
398    let n_feat = x_support[0].len();
399
400    let mut col_mean = vec![0f32; n_feat];
401    for f in 0..n_feat {
402        let mut sum = 0f64;
403        let mut count = 0usize;
404        for row in x_support {
405            if !row[f].is_nan() {
406                sum += row[f] as f64;
407                count += 1;
408            }
409        }
410        col_mean[f] = if count > 0 { (sum / count as f64) as f32 } else { 0.0 };
411    }
412
413    let impute = |row: &[f32]| -> Vec<f32> {
414        row.iter().enumerate().map(|(f, &v)| if v.is_nan() { col_mean[f] } else { v }).collect()
415    };
416    let support_imputed: Vec<Vec<f32>> = x_support.iter().map(|r| impute(r)).collect();
417    let query_imputed: Vec<Vec<f32>> = x_query.iter().map(|r| impute(r)).collect();
418
419    let kept: Vec<usize> = (0..n_feat)
420        .filter(|&f| {
421            let first = support_imputed[0][f];
422            support_imputed.iter().any(|r| r[f] != first)
423        })
424        .collect();
425
426    let select = |rows: &[Vec<f32>]| -> Vec<Vec<f32>> {
427        rows.iter().map(|r| kept.iter().map(|&f| r[f]).collect()).collect()
428    };
429
430    (select(&support_imputed), select(&query_imputed), kept)
431}
432
433/// Per-feature quantile-bucketize normalization (`Tab2DQuantileEmbeddingX`): fit 999 quantile
434/// boundaries per column from `x_support`, bucketize both `x_support`/`x_query` against them,
435/// normalize the bucket index by `n_support`, then z-score using `x_support`'s own mean/std.
436fn quantile_bucketize_normalize(x_support: &[Vec<f32>], x_query: &[Vec<f32>]) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
437    let n_support = x_support.len();
438    let n_features = x_support[0].len();
439    let n_query = x_query.len();
440
441    let mut out_support = vec![vec![0f32; n_features]; n_support];
442    let mut out_query = vec![vec![0f32; n_features]; n_query];
443
444    for feat in 0..n_features {
445        let mut col: Vec<f32> = x_support.iter().map(|r| r[feat]).collect();
446        col.sort_by(|a, b| a.partial_cmp(b).unwrap());
447
448        let boundaries: Vec<f32> = (1..1000).map(|i| quantile_linear(&col, i as f64 / 1000.0)).collect();
449
450        let bucket = |v: f32| -> f32 {
451            let idx = boundaries.partition_point(|&b| b < v);
452            idx as f32 / n_support as f32
453        };
454
455        let bucketed_support: Vec<f32> = x_support.iter().map(|r| bucket(r[feat])).collect();
456        let mean: f32 = bucketed_support.iter().sum::<f32>() / n_support as f32;
457        let var: f32 = bucketed_support.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / n_support as f32;
458        let std = var.sqrt();
459
460        for (i, &v) in bucketed_support.iter().enumerate() {
461            out_support[i][feat] = if std == 0.0 { 0.0 } else { (v - mean) / std };
462        }
463        for (i, row) in x_query.iter().enumerate() {
464            let v = bucket(row[feat]);
465            out_query[i][feat] = if std == 0.0 { 0.0 } else { (v - mean) / std };
466        }
467    }
468
469    (out_support, out_query)
470}
471
472/// `torch.quantile` default ('linear') interpolation on an already-sorted slice.
473fn quantile_linear(sorted: &[f32], q: f64) -> f32 {
474    let n = sorted.len();
475    if n == 1 {
476        return sorted[0];
477    }
478    let idx = q * (n - 1) as f64;
479    let lo = idx.floor() as usize;
480    let hi = idx.ceil() as usize;
481    let frac = (idx - lo as f64) as f32;
482    sorted[lo] + frac * (sorted[hi] - sorted[lo])
483}
zsfm_mitra/infer/mod.rs

zsfm_mitra/infer/
mod.rs
