1use std::path::Path;
9
10use anyhow::{Context, Result};
11use burn::tensor::{activation, Device, Tensor, TensorData};
12use zsfm_burn::linear::{burn_linear_bias, burn_linear_nobias};
13use zsfm_burn::norm::{burn_layer_norm_nd, burn_rms_norm_nd};
14use zsfm_burn::{open_gguf_weights, GgufWeights};
15
16use crate::config::TabDptConfig;
17
18const LN_EPS: f32 = 1e-5;
19const RMS_EPS: f32 = 1e-8;
20const CLIP_N_SIGMA: f32 = 8.0;
21
22fn device() -> Device {
23 Device::flex()
24}
25
26fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
27 weights.load_2d(name)
28}
29
30fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
31 weights.load_1d(name)
32}
33
34fn linear(x: Tensor<2>, w: Tensor<2>, b: Tensor<1>) -> Tensor<2> {
35 burn_linear_bias(x, w, b)
36}
37
38fn linear_o<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
39 zsfm_burn::linear::burn_linear_nd(x, w, b)
40}
41
42fn linear_nd<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
43 zsfm_burn::linear::burn_linear_nd(x, w, b)
44}
45
46fn layer_norm<const D: usize>(x: Tensor<D>, w: Tensor<1>, b: Tensor<1>) -> Tensor<D> {
47 burn_layer_norm_nd(x, w, b, LN_EPS)
48}
49
50fn layer_norm_no_affine(x: Tensor<2>) -> Tensor<2> {
51 zsfm_burn::norm::burn_layer_norm_no_affine(x, LN_EPS)
52}
53
54fn rms<const D: usize>(x: Tensor<D>, w: Tensor<1>) -> Tensor<D> {
55 burn_rms_norm_nd(x, w, RMS_EPS)
56}
57
58fn gelu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
59 activation::gelu(x)
60}
61
62fn silu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
63 activation::silu(x)
64}
65
66fn sigmoid<const D: usize>(x: Tensor<D>) -> Tensor<D> {
67 activation::sigmoid(x)
68}
69
70fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
71 activation::softmax(x, D - 1)
72}
73
74fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
75 t.to_data()
76 .try_to_vec()
77 .map_err(|e| anyhow::anyhow!("{e:?}"))
78}
79
80struct BurnBlockW {
81 attn_norm_w: Tensor<1>,
82 attn_norm_b: Tensor<1>,
83 ff_norm_w: Tensor<1>,
84 ff_norm_b: Tensor<1>,
85 q_proj_w: Tensor<2>,
86 k_proj_w: Tensor<2>,
87 v_proj_w: Tensor<2>,
88 out_proj_w: Tensor<2>,
89 q_gate_w: Tensor<2>,
90 q_norm_w: Tensor<1>,
91 k_norm_w: Tensor<1>,
92 ff_up_w: Tensor<2>,
93 ff_down_w: Tensor<2>,
94}
95
96struct BurnYEncW {
97 fc1_w: Tensor<2>,
98 fc1_b: Tensor<1>,
99 fc2_w: Tensor<2>,
100 fc2_b: Tensor<1>,
101}
102
103pub struct BurnTabDptModel {
104 config: TabDptConfig,
105 encoder_w: Tensor<2>,
106 encoder_b: Tensor<1>,
107 thinking_embed: Tensor<2>,
108 blocks: Vec<BurnBlockW>,
109 y_encoders: Vec<BurnYEncW>,
110 head_fc1_w: Tensor<2>,
111 head_fc1_b: Tensor<1>,
112 head_fc2_w: Tensor<2>,
113 head_fc2_b: Tensor<1>,
114}
115
116impl BurnTabDptModel {
117 pub fn load(gguf_path: &Path, config: TabDptConfig) -> Result<Self> {
118 let weights = open_gguf_weights(gguf_path)
119 .with_context(|| format!("open {}", gguf_path.display()))?;
120
121 let encoder_w = load_t(&weights, "encoder.weight")?;
122 let encoder_b = load_v(&weights, "encoder.bias")?;
123 let thinking_embed = load_t(&weights, "thinking_embed")?;
124
125 let mut blocks = Vec::with_capacity(config.n_layers);
126 let mut y_encoders = Vec::with_capacity(config.n_layers);
127 for n in 0..config.n_layers {
128 let p = |s: &str| format!("blk.{n}.{s}");
129 blocks.push(BurnBlockW {
130 attn_norm_w: load_v(&weights, &p("attn_norm.weight"))?,
131 attn_norm_b: load_v(&weights, &p("attn_norm.bias"))?,
132 ff_norm_w: load_v(&weights, &p("ff_norm.weight"))?,
133 ff_norm_b: load_v(&weights, &p("ff_norm.bias"))?,
134 q_proj_w: load_t(&weights, &p("q_proj.weight"))?,
135 k_proj_w: load_t(&weights, &p("k_proj.weight"))?,
136 v_proj_w: load_t(&weights, &p("v_proj.weight"))?,
137 out_proj_w: load_t(&weights, &p("out_proj.weight"))?,
138 q_gate_w: load_t(&weights, &p("q_gate.weight"))?,
139 q_norm_w: load_v(&weights, &p("q_norm.weight"))?,
140 k_norm_w: load_v(&weights, &p("k_norm.weight"))?,
141 ff_up_w: load_t(&weights, &p("ff_up.weight"))?,
142 ff_down_w: load_t(&weights, &p("ff_down.weight"))?,
143 });
144 let yp = |s: &str| format!("y_enc.{n}.{s}");
145 y_encoders.push(BurnYEncW {
146 fc1_w: load_t(&weights, &yp("fc1.weight"))?,
147 fc1_b: load_v(&weights, &yp("fc1.bias"))?,
148 fc2_w: load_t(&weights, &yp("fc2.weight"))?,
149 fc2_b: load_v(&weights, &yp("fc2.bias"))?,
150 });
151 }
152
153 Ok(Self {
154 config,
155 encoder_w,
156 encoder_b,
157 thinking_embed,
158 blocks,
159 y_encoders,
160 head_fc1_w: load_t(&weights, "head_fc1.weight")?,
161 head_fc1_b: load_v(&weights, "head_fc1.bias")?,
162 head_fc2_w: load_t(&weights, "head_fc2.weight")?,
163 head_fc2_b: load_v(&weights, "head_fc2.bias")?,
164 })
165 }
166
167 pub fn predict_classification(
168 &self,
169 x_support: &[Vec<f32>],
170 y_support: &[usize],
171 x_query: &[Vec<f32>],
172 n_classes: usize,
173 ) -> Result<Vec<Vec<f32>>> {
174 let (x_s, x_q) = preprocess_x_outer(x_support, x_query, self.config.max_num_features);
175 let y_s: Vec<f32> = y_support.iter().map(|&c| c as f32).collect();
176 let out = self.forward(&x_s, &y_s, &x_q)?;
177 let n_q = x_query.len();
178 let width = self.config.max_num_classes + self.config.regression_bin_count;
179 let flat: Vec<f32> = to_host(&out)?;
180 let mut result = Vec::with_capacity(n_q);
181 for i in 0..n_q {
182 let row = &flat[i * width..i * width + n_classes];
183 let log_probs = log_softmax(row);
184 result.push(softmax_host(&log_probs));
185 }
186 Ok(result)
187 }
188
189 pub fn predict_regression(
190 &self,
191 x_support: &[Vec<f32>],
192 y_support: &[f32],
193 x_query: &[Vec<f32>],
194 ) -> Result<Vec<f32>> {
195 let (x_s, x_q) = preprocess_x_outer(x_support, x_query, self.config.max_num_features);
196 let mean_y: f64 = y_support.iter().map(|&v| v as f64).sum::<f64>() / y_support.len() as f64;
197 let var_y: f64 = y_support
198 .iter()
199 .map(|&v| (v as f64 - mean_y).powi(2))
200 .sum::<f64>()
201 / (y_support.len() - 1) as f64;
202 let std_y = var_y.sqrt() + 1e-6;
203 let y_s: Vec<f32> = y_support
204 .iter()
205 .map(|&v| ((v as f64 - mean_y) / std_y) as f32)
206 .collect();
207 let out = self.forward(&x_s, &y_s, &x_q)?;
208 let n_q = x_query.len();
209 let width = self.config.max_num_classes + self.config.regression_bin_count;
210 let flat: Vec<f32> = to_host(&out)?;
211 let n_bins = self.config.regression_bin_count;
212 let bin_width =
213 (self.config.regression_bin_max - self.config.regression_bin_min) / n_bins as f32;
214 let bin_centers: Vec<f32> = (0..n_bins)
215 .map(|i| self.config.regression_bin_min + (i as f32 + 0.5) * bin_width)
216 .collect();
217 let mut result = Vec::with_capacity(n_q);
218 for i in 0..n_q {
219 let row = &flat[i * width + self.config.max_num_classes..i * width + width];
220 let probs = softmax_host(row);
221 let expectation: f32 = probs.iter().zip(&bin_centers).map(|(p, c)| p * c).sum();
222 result.push((expectation as f64 * std_y + mean_y) as f32);
223 }
224 Ok(result)
225 }
226
227 fn forward(
228 &self,
229 x_support: &[Vec<f32>],
230 y_support: &[f32],
231 x_query: &[Vec<f32>],
232 ) -> Result<Tensor<2>> {
233 let dev = device();
234 let n_s = x_support.len();
235 let n_q = x_query.len();
236 let n_feat = self.config.max_num_features;
237 let n_think = self.config.n_thinking_rows;
238 let ctx_len = n_think + n_s;
239
240 let mut all_rows: Vec<Vec<f32>> = Vec::with_capacity(n_s + n_q);
241 all_rows.extend_from_slice(x_support);
242 all_rows.extend_from_slice(x_query);
243 let all_rows = clip_outliers(all_rows, n_s, n_feat, CLIP_N_SIGMA);
244 let (all_rows, _m, _s) = normalize_data(all_rows, n_s, n_feat);
245 let mut all_rows = clip_outliers(all_rows, n_s, n_feat, CLIP_N_SIGMA);
246 for row in &mut all_rows {
247 for v in row.iter_mut() {
248 if v.is_nan() || v.is_infinite() {
249 *v = 0.0;
250 }
251 }
252 }
253 let flat: Vec<f32> = all_rows.into_iter().flatten().collect();
254 let x_t = Tensor::<2>::from_data(TensorData::new(flat, [n_s + n_q, n_feat]), &dev);
255 let x_enc = linear(x_t, self.encoder_w.clone(), self.encoder_b.clone());
256 let x_enc = layer_norm_no_affine(x_enc);
257 let mut src = Tensor::cat(vec![self.thinking_embed.clone(), x_enc], 0);
258
259 let y_support_t =
260 Tensor::<2>::from_data(TensorData::new(y_support.to_vec(), [n_s, 1]), &dev);
261 let zeros_think_emb = Tensor::<2>::zeros([n_think, self.config.y_encoder_dim], &dev);
262 let kappa = self.config.kappa();
263 let beta = attention_beta(kappa, ctx_len, self.config.base_len, self.config.max_len);
264 for (blk, yenc) in self.blocks.iter().zip(self.y_encoders.iter()) {
265 let y_h = gelu(linear(
266 y_support_t.clone(),
267 yenc.fc1_w.clone(),
268 yenc.fc1_b.clone(),
269 ));
270 let y_h = linear(y_h, yenc.fc2_w.clone(), yenc.fc2_b.clone());
271 let y_support_emb = layer_norm_no_affine(y_h);
272 let y_emb = Tensor::cat(vec![zeros_think_emb.clone(), y_support_emb], 0);
273 src = layer_forward(src, y_emb, ctx_len, blk, self.config.n_heads, beta as f32);
274 }
275 let query_out: Tensor<2> = src.narrow(0, ctx_len, n_q);
276 let h = gelu(linear(
277 query_out,
278 self.head_fc1_w.clone(),
279 self.head_fc1_b.clone(),
280 ));
281 Ok(linear_o(
282 h,
283 self.head_fc2_w.clone(),
284 Some(self.head_fc2_b.clone()),
285 ))
286 }
287}
288
289fn layer_forward(
290 src: Tensor<2>,
291 y_emb: Tensor<2>,
292 ctx_len: usize,
293 blk: &BurnBlockW,
294 n_heads: usize,
295 beta: f32,
296) -> Tensor<2> {
297 let dim = src.dims()[1];
298 let head_dim = dim / n_heads;
299 let l = src.dims()[0];
300 let h = layer_norm(
301 src.clone(),
302 blk.attn_norm_w.clone(),
303 blk.attn_norm_b.clone(),
304 );
305 let q = linear_nd(h.clone(), blk.q_proj_w.clone(), None);
306 let gate = sigmoid(linear_nd(q.clone(), blk.q_gate_w.clone(), None));
307 let h_ctx: Tensor<2> = h.narrow(0, 0, ctx_len);
308 let k = linear_nd(h_ctx.clone(), blk.k_proj_w.clone(), None);
309 let v_in = Tensor::cat(vec![h_ctx, y_emb], 1);
310 let v = linear_nd(v_in, blk.v_proj_w.clone(), None);
311 let q: Tensor<3> = q.reshape([l, n_heads, head_dim]).permute([1, 0, 2]);
312 let k: Tensor<3> = k.reshape([ctx_len, n_heads, head_dim]).permute([1, 0, 2]);
313 let v: Tensor<3> = v.reshape([ctx_len, n_heads, head_dim]).permute([1, 0, 2]);
314 let q = rms(q, blk.q_norm_w.clone());
315 let k = rms(k, blk.k_norm_w.clone());
316 let scale = beta / (head_dim as f32).sqrt();
317 let scores = q.matmul(k.transpose()).mul_scalar(scale);
318 let attn = softmax_last(scores);
319 let out = attn.matmul(v);
320 let out: Tensor<3> = out.permute([1, 0, 2]);
321 let gate: Tensor<3> = gate.reshape([l, n_heads, 1]);
322 let out = out * gate;
323 let out: Tensor<2> = out.reshape([l, dim]);
324 let attn_out = linear_nd(out, blk.out_proj_w.clone(), None);
325 let x1 = src + attn_out;
326 let ff_in = layer_norm(x1.clone(), blk.ff_norm_w.clone(), blk.ff_norm_b.clone());
327 let up = linear_nd(ff_in, blk.ff_up_w.clone(), None);
328 let ff_dim2 = up.dims()[1] / 2;
329 let u: Tensor<2> = up.clone().narrow(1, 0, ff_dim2);
330 let v_ff: Tensor<2> = up.narrow(1, ff_dim2, ff_dim2);
331 let ff_out = linear_nd(silu(u) * v_ff, blk.ff_down_w.clone(), None);
332 x1 + ff_out
333}
334
335fn attention_beta(kappa: Option<f64>, ctx_len: usize, base_len: usize, max_len: usize) -> f64 {
336 match kappa {
337 None => 1.0,
338 Some(kappa) => {
339 let n = (ctx_len as f64).min(max_len as f64);
340 let log_term = (n / base_len as f64).ln().max(0.0);
341 1.0 + kappa * log_term
342 }
343 }
344}
345
346fn softmax_host(logits: &[f32]) -> Vec<f32> {
347 zsfm_core::softmax(logits)
348}
349
350fn log_softmax(logits: &[f32]) -> Vec<f32> {
351 let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
352 let log_sum_exp = logits.iter().map(|&v| (v - max).exp()).sum::<f32>().ln();
353 logits.iter().map(|&v| v - max - log_sum_exp).collect()
354}
355
356fn preprocess_x_outer(
357 x_support: &[Vec<f32>],
358 x_query: &[Vec<f32>],
359 max_features: usize,
360) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
361 let n_feat = x_support[0].len();
362 let mut impute_mean = vec![0f32; n_feat];
363 for f in 0..n_feat {
364 let mut sum = 0f64;
365 let mut count = 0usize;
366 for row in x_support {
367 if !row[f].is_nan() {
368 sum += row[f] as f64;
369 count += 1;
370 }
371 }
372 impute_mean[f] = if count > 0 {
373 (sum / count as f64) as f32
374 } else {
375 0.0
376 };
377 }
378 let impute = |row: &[f32]| -> Vec<f32> {
379 row.iter()
380 .enumerate()
381 .map(|(f, &v)| if v.is_nan() { impute_mean[f] } else { v })
382 .collect()
383 };
384 let support_imputed: Vec<Vec<f32>> = x_support.iter().map(|r| impute(r)).collect();
385 let query_imputed: Vec<Vec<f32>> = x_query.iter().map(|r| impute(r)).collect();
386 let n_s = support_imputed.len();
387 let mut scaler_mean = vec![0f32; n_feat];
388 let mut scaler_std = vec![0f32; n_feat];
389 for f in 0..n_feat {
390 let mean: f64 = support_imputed.iter().map(|r| r[f] as f64).sum::<f64>() / n_s as f64;
391 let var: f64 = support_imputed
392 .iter()
393 .map(|r| (r[f] as f64 - mean).powi(2))
394 .sum::<f64>()
395 / n_s as f64;
396 scaler_mean[f] = mean as f32;
397 scaler_std[f] = if var == 0.0 { 1.0 } else { var.sqrt() as f32 };
398 }
399 let scale_and_pad = |rows: &[Vec<f32>]| -> Vec<Vec<f32>> {
400 rows.iter()
401 .map(|row| {
402 let mut out = vec![0f32; max_features];
403 for f in 0..n_feat.min(max_features) {
404 out[f] = (row[f] - scaler_mean[f]) / scaler_std[f];
405 }
406 out
407 })
408 .collect()
409 };
410 (
411 scale_and_pad(&support_imputed),
412 scale_and_pad(&query_imputed),
413 )
414}
415
416fn clip_outliers(
417 mut rows: Vec<Vec<f32>>,
418 n_ctx: usize,
419 n_feat: usize,
420 n_sigma: f32,
421) -> Vec<Vec<f32>> {
422 for f in 0..n_feat {
423 let ctx_vals: Vec<f32> = rows[..n_ctx]
424 .iter()
425 .map(|r| r[f])
426 .filter(|v| !v.is_nan())
427 .collect();
428 if ctx_vals.is_empty() {
429 continue;
430 }
431 let mean = ctx_vals.iter().sum::<f32>() / ctx_vals.len() as f32;
432 let std1 = population_std(&ctx_vals, mean);
433 let cutoff1 = n_sigma * std1;
434 let refined: Vec<f32> = ctx_vals
435 .iter()
436 .copied()
437 .filter(|&v| (v - mean).abs() <= cutoff1)
438 .collect();
439 let std2 = if refined.len() < 2 {
440 std1
441 } else {
442 let refined_mean = refined.iter().sum::<f32>() / refined.len() as f32;
443 population_std(&refined, refined_mean)
444 };
445 let cutoff2 = n_sigma * std2;
446 let (lo, hi) = (mean - cutoff2, mean + cutoff2);
447 for row in rows.iter_mut() {
448 row[f] = row[f].clamp(lo, hi);
449 }
450 }
451 rows
452}
453
454fn normalize_data(
455 rows: Vec<Vec<f32>>,
456 n_ctx: usize,
457 n_feat: usize,
458) -> (Vec<Vec<f32>>, Vec<f32>, Vec<f32>) {
459 let mut mean = vec![0f32; n_feat];
460 let mut std = vec![0f32; n_feat];
461 for f in 0..n_feat {
462 let ctx_vals: Vec<f32> = rows[..n_ctx]
463 .iter()
464 .map(|r| r[f])
465 .filter(|v| !v.is_nan())
466 .collect();
467 if ctx_vals.is_empty() {
468 mean[f] = 0.0;
469 std[f] = 1e-6;
470 continue;
471 }
472 let m = ctx_vals.iter().sum::<f32>() / ctx_vals.len() as f32;
473 mean[f] = m;
474 std[f] = population_std(&ctx_vals, m) + 1e-6;
475 }
476 let out: Vec<Vec<f32>> = rows
477 .into_iter()
478 .map(|row| {
479 row.iter()
480 .enumerate()
481 .map(|(f, &v)| (v - mean[f]) / std[f])
482 .collect()
483 })
484 .collect();
485 (out, mean, std)
486}
487
488fn population_std(vals: &[f32], mean: f32) -> f32 {
489 if vals.len() < 2 {
490 return 0.0;
491 }
492 let var = vals.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / (vals.len() - 1) as f32;
493 var.sqrt()
494}
495
496#[cfg(test)]
497mod tests {
498 use super::*;
499 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
500
501 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
502 let mut x = seed;
503 (0..n)
504 .map(|_| {
505 x = x
506 .wrapping_mul(6364136223846793005)
507 .wrapping_add(1442695040888963407);
508 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
509 })
510 .collect()
511 }
512
513 fn tiny_config() -> TabDptConfig {
514 TabDptConfig {
515 dim: 8,
516 n_layers: 1,
517 n_heads: 2,
518 ff_dim: 8,
519 y_encoder_dim: 4,
520 max_num_classes: 2,
521 regression_bin_count: 4,
522 regression_bin_min: -2.0,
523 regression_bin_max: 2.0,
524 max_num_features: 4,
525 base_len: 64,
526 max_len: 64,
527 n_thinking_rows: 2,
528 }
529 }
530
531 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
532 *seed += 1;
533 let data: Vec<f32> = pseudo(rows * cols, *seed)
534 .iter()
535 .map(|v| v * scale)
536 .collect();
537 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
538 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
539 }
540
541 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
542 *seed += 1;
543 let bytes: Vec<u8> = pseudo(n, *seed)
544 .iter()
545 .flat_map(|v| v.to_le_bytes())
546 .collect();
547 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
548 }
549
550 #[test]
551 fn burn_classification_golden() {
552 let dir = std::env::temp_dir().join(format!("tabdpt-burn-{}", std::process::id()));
553 let _ = std::fs::remove_dir_all(&dir);
554 std::fs::create_dir_all(&dir).unwrap();
555 let path = dir.join("model.gguf");
556 let mut w = GGUFWriter::new();
557 w.add_metadata(
558 "general.architecture",
559 GGUFMetaValue::String("tabdpt".into()),
560 );
561 let mut seed = 1900u64;
562 let s = 0.2;
563 put2(&mut w, "encoder.weight", 8, 4, &mut seed, s);
564 put1(&mut w, "encoder.bias", 8, &mut seed);
565 put2(&mut w, "thinking_embed", 2, 8, &mut seed, s);
566 put1(&mut w, "blk.0.attn_norm.weight", 8, &mut seed);
567 put1(&mut w, "blk.0.attn_norm.bias", 8, &mut seed);
568 put1(&mut w, "blk.0.ff_norm.weight", 8, &mut seed);
569 put1(&mut w, "blk.0.ff_norm.bias", 8, &mut seed);
570 put2(&mut w, "blk.0.q_proj.weight", 8, 8, &mut seed, s);
571 put2(&mut w, "blk.0.k_proj.weight", 8, 8, &mut seed, s);
572 put2(&mut w, "blk.0.v_proj.weight", 8, 12, &mut seed, s);
573 put2(&mut w, "blk.0.out_proj.weight", 8, 8, &mut seed, s);
574 put2(&mut w, "blk.0.q_gate.weight", 2, 8, &mut seed, s);
575 put1(&mut w, "blk.0.q_norm.weight", 4, &mut seed);
576 put1(&mut w, "blk.0.k_norm.weight", 4, &mut seed);
577 put2(&mut w, "blk.0.ff_up.weight", 16, 8, &mut seed, s);
578 put2(&mut w, "blk.0.ff_down.weight", 8, 8, &mut seed, s);
579 put2(&mut w, "y_enc.0.fc1.weight", 4, 1, &mut seed, s);
580 put1(&mut w, "y_enc.0.fc1.bias", 4, &mut seed);
581 put2(&mut w, "y_enc.0.fc2.weight", 4, 4, &mut seed, s);
582 put1(&mut w, "y_enc.0.fc2.bias", 4, &mut seed);
583 put2(&mut w, "head_fc1.weight", 8, 8, &mut seed, s);
584 put1(&mut w, "head_fc1.bias", 8, &mut seed);
585 put2(&mut w, "head_fc2.weight", 6, 8, &mut seed, s);
586 put1(&mut w, "head_fc2.bias", 6, &mut seed);
587 let f = std::fs::File::create(&path).unwrap();
588 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
589
590 let cfg = tiny_config();
591 let burn = BurnTabDptModel::load(&path, cfg).unwrap();
592 let x_support = vec![
593 vec![0.1, 1.2],
594 vec![0.9, -0.3],
595 vec![-1.1, 0.4],
596 vec![1.5, 1.1],
597 ];
598 let y_support = vec![0usize, 1, 1, 0];
599 let x_query = vec![vec![0.2, 0.9], vec![-0.8, 0.1]];
600 let t0 = std::time::Instant::now();
601 let a = burn
602 .predict_classification(&x_support, &y_support, &x_query, 2)
603 .unwrap();
604 let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
605 let b = burn
606 .predict_classification(&x_support, &y_support, &x_query, 2)
607 .unwrap();
608 assert_eq!(a.len(), 2);
609 assert!(a
610 .iter()
611 .all(|r| r.len() == 2 && r.iter().all(|v| v.is_finite())));
612 assert_eq!(a, b, "predict must be deterministic");
613 println!("tabdpt synthetic: burn {burn_ms:.1}ms out={a:?}");
614 let golden = [[0.5918571, 0.4081429], [0.5925009, 0.40749905]];
616 for (ra, re) in a.iter().zip(golden.iter()) {
617 for (x, y) in ra.iter().zip(re.iter()) {
618 assert!((x - y).abs() < 1e-4, "golden drift: {x} vs {y}");
619 }
620 }
621 let _ = std::fs::remove_dir_all(&dir);
622 }
623}