1use std::path::Path;
12
13use anyhow::{Context, Result};
14use burn::tensor::{activation, Device, Tensor, TensorData};
15use zsfm_burn::linear::burn_linear_nd;
16use zsfm_burn::norm::burn_rms_norm_nd;
17use zsfm_burn::{open_gguf_weights, GgufWeights};
18
19use crate::config::TabFMConfig;
20
21fn device() -> Device {
22 Device::flex()
23}
24
25fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
26 weights.load_2d(name)
27}
28
29fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
30 weights.load_1d(name)
31}
32
33fn load_w(weights: &GgufWeights, name: &str, expected_d_out: usize) -> Result<Tensor<2>> {
34 weights.load_weight_2d(name, expected_d_out)
35}
36
37fn load_flat(weights: &GgufWeights, name: &str) -> Result<Vec<f32>> {
38 weights.load_flat(name)
39}
40
41fn linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
42 burn_linear_nd(x, w, b)
43}
44
45fn rms<const D: usize>(x: Tensor<D>, w: Tensor<1>, eps: f64) -> Tensor<D> {
46 burn_rms_norm_nd(x, w, eps as f32)
47}
48
49fn sigmoid<const D: usize>(x: Tensor<D>) -> Tensor<D> {
50 (x.neg().exp() + 1.0).recip()
51}
52
53fn silu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
54 x.clone() * sigmoid(x)
55}
56
57fn softplus<const D: usize>(x: Tensor<D>) -> Tensor<D> {
58 (x.exp() + 1.0).log()
59}
60
61fn gelu_tanh<const D: usize>(x: Tensor<D>) -> Tensor<D> {
63 activation::gelu_approximate(x)
64}
65
66fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
67 activation::softmax(x, D - 1)
68}
69
70fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
71 t.to_data()
72 .try_to_vec()
73 .map_err(|e| anyhow::anyhow!("{e:?}"))
74}
75
76fn to_host_1(t: &Tensor<1>) -> Result<Vec<f32>> {
77 to_host(t)
78}
79
80#[derive(Clone)]
81struct BurnMabW {
82 q_w: Tensor<2>,
83 q_b: Tensor<1>,
84 k_w: Tensor<2>,
85 k_b: Tensor<1>,
86 v_w: Tensor<2>,
87 v_b: Tensor<1>,
88 o_w: Tensor<2>,
89 o_b: Tensor<1>,
90 q_norm: Tensor<1>,
91 k_norm: Tensor<1>,
92 per_dim_scale: Tensor<1>,
93 pre_attn_norm: Tensor<1>,
94 post_attn_norm: Tensor<1>,
95 pre_ff_norm: Tensor<1>,
96 post_ff_norm: Tensor<1>,
97 ffn_up_w: Tensor<2>,
98 ffn_up_b: Tensor<1>,
99 ffn_gate_w: Tensor<2>,
100 ffn_gate_b: Tensor<1>,
101 ffn_down_w: Tensor<2>,
102 ffn_down_b: Tensor<1>,
103}
104
105struct BurnInducedBlockW {
106 ind_vectors: Tensor<2>,
107 mab1: BurnMabW,
108 mab2: BurnMabW,
109}
110
111struct BurnColStackW {
112 blocks: Vec<BurnInducedBlockW>,
113 out_w_w: Tensor<2>,
114 out_w_b: Tensor<1>,
115 out_norm_w: Tensor<1>,
116}
117
118struct BurnRowStackW {
119 rope_freqs: Vec<f32>,
120 blocks: Vec<BurnMabW>,
121 out_norm_w: Tensor<1>,
122}
123
124struct BurnMlpW {
125 layers: Vec<(Tensor<2>, Tensor<1>)>,
126}
127
128enum BurnYEmbedW {
129 Embedding(Tensor<2>),
130 Mlp(BurnMlpW),
131}
132
133enum BurnYEncoderW {
134 OneHot {
135 proj_w: Tensor<2>,
136 proj_b: Tensor<1>,
137 },
138 Mlp(BurnMlpW),
139}
140
141struct BurnCellW {
142 fourier_freq: Tensor<2>,
143 fourier_freq_cat: Tensor<2>,
144 in_linear_w: Tensor<2>,
145 in_linear_b: Tensor<1>,
146 in_linear_cat_w: Tensor<2>,
147 in_linear_cat_b: Tensor<1>,
148 y_embed: BurnYEmbedW,
149}
150
151struct BurnIclW {
152 blocks: Vec<BurnMabW>,
153 out_norm_w: Tensor<1>,
154 y_encoder: BurnYEncoderW,
155 decoder: BurnMlpW,
156}
157
158pub struct BurnTabFMModel {
159 config: BurnInferConfig,
160 cell: BurnCellW,
161 colenc1: BurnColStackW,
162 colenc2: BurnColStackW,
163 rowenc1: BurnRowStackW,
164 rowenc2: BurnRowStackW,
165 cls_tokens: Tensor<2>,
166 icl: BurnIclW,
167}
168
169#[derive(Clone)]
170struct BurnInferConfig {
171 embed_dim: usize,
172 max_classes: usize,
173 col_nhead: usize,
174 row_nhead: usize,
175 row_num_cls: usize,
176 icl_nhead: usize,
177 ff_factor: usize,
178 feature_group_size: usize,
179 num_freq: usize,
180 norm_eps: f64,
181 is_classifier: bool,
182}
183
184fn load_mab(weights: &GgufWeights, prefix: &str, e: usize, ff: usize) -> Result<BurnMabW> {
185 let p = |s: &str| format!("{prefix}.{s}");
186 Ok(BurnMabW {
187 q_w: load_w(weights, &p("attn_q.weight"), e)?,
188 q_b: load_v(weights, &p("attn_q.bias"))?,
189 k_w: load_w(weights, &p("attn_k.weight"), e)?,
190 k_b: load_v(weights, &p("attn_k.bias"))?,
191 v_w: load_w(weights, &p("attn_v.weight"), e)?,
192 v_b: load_v(weights, &p("attn_v.bias"))?,
193 o_w: load_w(weights, &p("attn_o.weight"), e)?,
194 o_b: load_v(weights, &p("attn_o.bias"))?,
195 q_norm: load_v(weights, &p("q_norm.weight"))?,
196 k_norm: load_v(weights, &p("k_norm.weight"))?,
197 per_dim_scale: load_v(weights, &p("per_dim_scale"))?,
198 pre_attn_norm: load_v(weights, &p("pre_attn_norm.weight"))?,
199 post_attn_norm: load_v(weights, &p("post_attn_norm.weight"))?,
200 pre_ff_norm: load_v(weights, &p("pre_ff_norm.weight"))?,
201 post_ff_norm: load_v(weights, &p("post_ff_norm.weight"))?,
202 ffn_up_w: load_w(weights, &p("ffn_up.weight"), ff)?,
203 ffn_up_b: load_v(weights, &p("ffn_up.bias"))?,
204 ffn_gate_w: load_w(weights, &p("ffn_gate.weight"), ff)?,
205 ffn_gate_b: load_v(weights, &p("ffn_gate.bias"))?,
206 ffn_down_w: load_w(weights, &p("ffn_down.weight"), e)?,
207 ffn_down_b: load_v(weights, &p("ffn_down.bias"))?,
208 })
209}
210
211fn load_col_stack(
212 weights: &GgufWeights,
213 stack: &str,
214 num_blocks: usize,
215 e: usize,
216 ff: usize,
217) -> Result<BurnColStackW> {
218 let mut blocks = Vec::with_capacity(num_blocks);
219 for n in 0..num_blocks {
220 let ind_vectors = load_t(weights, &format!("{stack}.blk.{n}.ind_vectors"))?;
221 let mab1 = load_mab(weights, &format!("{stack}.blk.{n}.mab1"), e, ff)?;
222 let mab2 = load_mab(weights, &format!("{stack}.blk.{n}.mab2"), e, ff)?;
223 blocks.push(BurnInducedBlockW {
224 ind_vectors,
225 mab1,
226 mab2,
227 });
228 }
229 Ok(BurnColStackW {
230 blocks,
231 out_w_w: load_w(weights, &format!("{stack}.out_w.weight"), e)?,
232 out_w_b: load_v(weights, &format!("{stack}.out_w.bias"))?,
233 out_norm_w: load_v(weights, &format!("{stack}.out_norm.weight"))?,
234 })
235}
236
237fn load_row_stack(
238 weights: &GgufWeights,
239 stack: &str,
240 num_blocks: usize,
241 e: usize,
242 ff: usize,
243) -> Result<BurnRowStackW> {
244 let rope_freqs: Vec<f32> = load_flat(weights, &format!("{stack}.rope_freqs"))?;
245 let mut blocks = Vec::with_capacity(num_blocks);
246 for n in 0..num_blocks {
247 blocks.push(load_mab(weights, &format!("{stack}.blk.{n}"), e, ff)?);
248 }
249 Ok(BurnRowStackW {
250 rope_freqs,
251 blocks,
252 out_norm_w: load_v(weights, &format!("{stack}.out_norm.weight"))?,
253 })
254}
255
256fn load_mlp(weights: &GgufWeights, prefix: &str, dims: &[usize]) -> Result<BurnMlpW> {
257 let mut layers = Vec::with_capacity(dims.len() - 1);
258 for i in 0..dims.len() - 1 {
259 let out_dim = dims[i + 1];
260 let w = load_w(weights, &format!("{prefix}.mlp.{i}.weight"), out_dim)?;
261 let b = load_v(weights, &format!("{prefix}.mlp.{i}.bias"))?;
262 layers.push((w, b));
263 }
264 Ok(BurnMlpW { layers })
265}
266
267struct BurnTabfmRope {
268 freqs: Vec<f32>,
269}
270
271impl BurnTabfmRope {
272 fn rotate(&self, x: Tensor<4>) -> Result<Tensor<4>> {
275 let xd = x.dims();
276 let (n, t, nh, hd) = (xd[0], xd[1], xd[2], xd[3]);
277 let dev = device();
278 let half = hd / 2;
279 let positions: Vec<f32> = (0..t).map(|p| p as f32).collect();
280 let pos = Tensor::<1>::from_data(TensorData::new(positions, [t]), &dev);
281 let freqs = Tensor::<1>::from_data(TensorData::new(self.freqs.clone(), [half]), &dev);
282 let f = pos.unsqueeze_dim::<2>(1) * freqs.unsqueeze_dim::<2>(0);
283 let cos = f.clone().cos();
284 let sin = f.sin();
285 let cos_i: Tensor<3> = Tensor::stack(vec![cos.clone(), cos], 2);
286 let cos_i: Tensor<2> = cos_i.reshape([t, hd]);
287 let sin_i: Tensor<3> = Tensor::stack(vec![sin.clone(), sin], 2);
288 let sin_i: Tensor<2> = sin_i.reshape([t, hd]);
289 let cos_i = cos_i.reshape([1, t, 1, hd]);
290 let sin_i = sin_i.reshape([1, t, 1, hd]);
291 let pairs = x.clone().reshape([n, t, nh, half, 2]);
292 let x1 = pairs.clone().narrow(4, 0, 1).squeeze_dim::<4>(4);
293 let x2 = pairs.narrow(4, 1, 1).squeeze_dim::<4>(4);
294 let rot_stack: Tensor<5> = Tensor::stack(vec![x2.mul_scalar(-1.0), x1], 4);
295 let rot = rot_stack.reshape([n, t, nh, hd]);
296 Ok(x * cos_i + rot * sin_i)
297 }
298}
299impl BurnTabFMModel {
300 pub fn load(gguf_path: &Path, tc: &TabFMConfig) -> Result<Self> {
301 let weights = open_gguf_weights(gguf_path)
302 .with_context(|| format!("open {}", gguf_path.display()))?;
303 let e = tc.embed_dim as usize;
304 let col_ff = e * tc.ff_factor as usize;
305 let icl_dim = e * tc.row_num_cls as usize;
306 let icl_ff = icl_dim * tc.ff_factor as usize;
307 let decoder_hidden = tc.decoder_hidden.map(|v| v as usize).unwrap_or(icl_dim * 2);
308 let out_dim = if tc.is_classifier {
309 tc.max_classes as usize
310 } else {
311 1
312 };
313 let config = BurnInferConfig {
314 embed_dim: e,
315 max_classes: tc.max_classes as usize,
316 col_nhead: tc.col_nhead as usize,
317 row_nhead: tc.row_nhead as usize,
318 row_num_cls: tc.row_num_cls as usize,
319 icl_nhead: tc.icl_nhead as usize,
320 ff_factor: tc.ff_factor as usize,
321 feature_group_size: tc.feature_group_size as usize,
322 num_freq: tc.num_freq as usize,
323 norm_eps: tc.norm_eps,
324 is_classifier: tc.is_classifier,
325 };
326 let cell = BurnCellW {
327 fourier_freq: load_t(&weights, "cell.fourier_freq")?,
328 fourier_freq_cat: load_t(&weights, "cell.fourier_freq_cat")?,
329 in_linear_w: load_w(&weights, "cell.in_linear.weight", e)?,
330 in_linear_b: load_v(&weights, "cell.in_linear.bias")?,
331 in_linear_cat_w: load_w(&weights, "cell.in_linear_cat.weight", e)?,
332 in_linear_cat_b: load_v(&weights, "cell.in_linear_cat.bias")?,
333 y_embed: if tc.is_classifier {
334 BurnYEmbedW::Embedding(load_t(&weights, "cell.y_embed.weight")?)
335 } else {
336 BurnYEmbedW::Mlp(load_mlp(&weights, "cell.y_embed", &[1, 6, e])?)
337 },
338 };
339 let colenc1 = load_col_stack(&weights, "colenc1", tc.col_num_blocks as usize, e, col_ff)?;
340 let colenc2 = load_col_stack(&weights, "colenc2", tc.col_num_blocks as usize, e, col_ff)?;
341 let rowenc1 = load_row_stack(&weights, "rowenc1", tc.row_num_blocks as usize, e, col_ff)?;
342 let rowenc2 = load_row_stack(&weights, "rowenc2", tc.row_num_blocks as usize, e, col_ff)?;
343 let cls_tokens = load_t(&weights, "cls_tokens")?;
344 let mut icl_blocks = Vec::with_capacity(tc.icl_num_blocks as usize);
345 for n in 0..tc.icl_num_blocks as usize {
346 icl_blocks.push(load_mab(
347 &weights,
348 &format!("icl.blk.{n}"),
349 icl_dim,
350 icl_ff,
351 )?);
352 }
353 let y_encoder = if tc.is_classifier {
354 BurnYEncoderW::OneHot {
355 proj_w: load_w(&weights, "icl.y_encoder.projection.weight", icl_dim)?,
356 proj_b: load_v(&weights, "icl.y_encoder.projection.bias")?,
357 }
358 } else {
359 BurnYEncoderW::Mlp(load_mlp(
360 &weights,
361 "icl.y_encoder",
362 &[1, decoder_hidden, icl_dim],
363 )?)
364 };
365 let decoder = load_mlp(&weights, "icl.decoder", &[icl_dim, decoder_hidden, out_dim])?;
366 let icl = BurnIclW {
367 blocks: icl_blocks,
368 out_norm_w: load_v(&weights, "icl.out_norm.weight")?,
369 y_encoder,
370 decoder,
371 };
372 Ok(Self {
373 config,
374 cell,
375 colenc1,
376 colenc2,
377 rowenc1,
378 rowenc2,
379 cls_tokens,
380 icl,
381 })
382 }
383
384 pub fn is_classifier(&self) -> bool {
385 self.config.is_classifier
386 }
387
388 pub fn predict(
389 &self,
390 x: &[Vec<f32>],
391 y: &[f32],
392 train_size: usize,
393 cat_mask: Option<&[bool]>,
394 d: Option<usize>,
395 ) -> Result<Vec<Vec<f32>>> {
396 let out = self.predict_batch(
397 &[x.to_vec()],
398 &[y.to_vec()],
399 train_size,
400 &[cat_mask.unwrap_or(&[]).to_vec()],
401 d,
402 )?;
403 out.into_iter()
404 .next()
405 .context("predict_batch returned no batch items")
406 }
407
408 pub fn predict_batch(
409 &self,
410 x_batch: &[Vec<Vec<f32>>],
411 y_batch: &[Vec<f32>],
412 train_size: usize,
413 cat_mask_batch: &[Vec<bool>],
414 d: Option<usize>,
415 ) -> Result<Vec<Vec<Vec<f32>>>> {
416 let cfg = &self.config;
417 let b_len = x_batch.len();
418 anyhow::ensure!(b_len > 0, "x_batch must have at least one table");
419 let t_len = x_batch[0].len();
420 let h_len = x_batch[0][0].len();
421 let d_val = d.unwrap_or(h_len).min(h_len);
422 let emb0 = self.cell_embed(x_batch, y_batch, train_size, cat_mask_batch, d_val)?;
423 let emb1 = self.col_embedding_forward(&emb0, train_size, &self.colenc1)?;
424 let num_cls = cfg.row_num_cls;
425 let cls: Tensor<4> = self
426 .cls_tokens
427 .clone()
428 .reshape([1, 1, num_cls, cfg.embed_dim])
429 .expand([b_len, t_len, num_cls, cfg.embed_dim]);
430 let emb2 = Tensor::cat(vec![cls, emb1], 2);
431 let d_plus_cls = d_val + num_cls;
432 let emb3 = self.row_interaction_forward(&emb2, d_plus_cls, &self.rowenc1)?;
433 let emb4 = self.col_embedding_forward(&emb3, train_size, &self.colenc2)?;
434 let reps = self.row_interaction_collapsed(&emb4, &self.rowenc2)?;
435 let logits = self.icl_forward(&reps, y_batch, train_size)?;
436 let out_dim = if cfg.is_classifier {
437 cfg.max_classes
438 } else {
439 1
440 };
441 let flat: Vec<f32> = to_host(&logits)?;
442 let mut result = vec![vec![vec![0f32; out_dim]; t_len]; b_len];
443 for (bb, batch_item) in result.iter_mut().enumerate() {
444 for (t, row) in batch_item.iter_mut().enumerate() {
445 let base = ((bb * t_len) + t) * out_dim;
446 row.copy_from_slice(&flat[base..base + out_dim]);
447 }
448 }
449 Ok(result)
450 }
451
452 fn cell_embed(
453 &self,
454 x_batch: &[Vec<Vec<f32>>],
455 y_batch: &[Vec<f32>],
456 train_size: usize,
457 cat_mask_batch: &[Vec<bool>],
458 d: usize,
459 ) -> Result<Tensor<4>> {
460 let cfg = &self.config;
461 let b_len = x_batch.len();
462 let t_len = x_batch[0].len();
463 let h_len = x_batch[0][0].len();
464 let fgs = cfg.feature_group_size;
465 let num_freq = cfg.num_freq;
466 let e = cfg.embed_dim;
467 let d_safe = d.max(1);
468 let mut idx = vec![vec![0usize; h_len]; fgs];
469 for (g, row) in idx.iter_mut().enumerate() {
470 let offset = (1usize << g) - 1;
471 for (h, slot) in row.iter_mut().enumerate() {
472 *slot = (h + offset) % d_safe;
473 }
474 }
475 let ff: Vec<f32> = to_host(&self.cell.fourier_freq)?;
476 let ffc: Vec<f32> = to_host(&self.cell.fourier_freq_cat)?;
477 let in_w: Vec<f32> = to_host(&self.cell.in_linear_w)?;
478 let in_b: Vec<f32> = to_host_1(&self.cell.in_linear_b)?;
479 let in_w_cat: Vec<f32> = to_host(&self.cell.in_linear_cat_w)?;
480 let in_b_cat: Vec<f32> = to_host_1(&self.cell.in_linear_cat_b)?;
481 let y_emb_t = compute_y_embed(y_batch, &self.cell.y_embed, cfg.max_classes)?;
482 let y_emb: Vec<f32> = to_host(&y_emb_t)?;
483 let mut out = vec![0f32; b_len * t_len * h_len * e];
484 for bb in 0..b_len {
485 let x = &x_batch[bb];
486 let cat_mask = &cat_mask_batch[bb];
487 for tt in 0..t_len {
488 let add_y = tt < train_size;
489 for hh in 0..h_len {
490 if hh >= d {
491 continue;
492 }
493 let mut acc = vec![0f32; e];
494 for g in 0..fgs {
495 let src_col = idx[g][hh];
496 let val = x[tt][src_col];
497 let is_cat = cat_mask.get(src_col).copied().unwrap_or(false);
498 let (freq_row, w_lin, b_lin): (&[f32], &[f32], &[f32]) = if is_cat {
499 (&ffc[g * num_freq..(g + 1) * num_freq], &in_w_cat, &in_b_cat)
500 } else {
501 (&ff[g * num_freq..(g + 1) * num_freq], &in_w, &in_b)
502 };
503 for ee in 0..e {
504 let mut s = b_lin[ee];
505 let row_off = ee * 2 * num_freq;
506 for f in 0..num_freq {
507 let arg = val * freq_row[f];
508 s += arg.sin() * w_lin[row_off + f];
509 s += arg.cos() * w_lin[row_off + num_freq + f];
510 }
511 acc[ee] += s;
512 }
513 }
514 let base = ((bb * t_len + tt) * h_len + hh) * e;
515 let y_base = (bb * t_len + tt) * e;
516 for ee in 0..e {
517 out[base + ee] = acc[ee] + if add_y { y_emb[y_base + ee] } else { 0.0 };
518 }
519 }
520 }
521 }
522 let dev = device();
523 Ok(Tensor::<4>::from_data(
524 TensorData::new(out, [b_len, t_len, h_len, e]),
525 &dev,
526 ))
527 }
528
529 fn col_embedding_forward(
530 &self,
531 x: &Tensor<4>,
532 train_size: usize,
533 w: &BurnColStackW,
534 ) -> Result<Tensor<4>> {
535 let cfg = &self.config;
536 let xd = x.dims();
537 let (b_len, t_len, hc, e) = (xd[0], xd[1], xd[2], xd[3]);
538 let src: Tensor<3> = x
539 .clone()
540 .permute([0, 2, 1, 3])
541 .reshape([b_len * hc, t_len, e]);
542 let mask = additive_key_mask(t_len, train_size);
543 let mut cur = src;
544 for blk in &w.blocks {
545 let n = cur.dims()[0];
546 let ind: Tensor<3> = blk.ind_vectors.clone().unsqueeze_dim::<3>(0).expand([
547 n,
548 blk.ind_vectors.dims()[0],
549 blk.ind_vectors.dims()[1],
550 ]);
551 let hidden = mab_forward(
552 &ind,
553 &cur,
554 &cur,
555 &blk.mab1,
556 cfg.col_nhead,
557 cfg.norm_eps,
558 Some(mask.clone()),
559 None,
560 )?;
561 cur = mab_forward(
562 &cur,
563 &hidden,
564 &hidden,
565 &blk.mab2,
566 cfg.col_nhead,
567 cfg.norm_eps,
568 None,
569 None,
570 )?;
571 }
572 let projected = linear(cur, w.out_w_w.clone(), Some(w.out_w_b.clone()));
573 let normed = rms(projected, w.out_norm_w.clone(), cfg.norm_eps);
574 Ok(normed.reshape([b_len, hc, t_len, e]).permute([0, 2, 1, 3]))
575 }
576
577 fn row_interaction_forward(
578 &self,
579 x: &Tensor<4>,
580 d_plus_cls: usize,
581 w: &BurnRowStackW,
582 ) -> Result<Tensor<4>> {
583 let out3 = self.row_interaction_3d(x, d_plus_cls, w)?;
584 let xd = x.dims();
585 Ok(out3.reshape([xd[0], xd[1], xd[2], xd[3]]))
586 }
587
588 fn row_interaction_collapsed(&self, x: &Tensor<4>, w: &BurnRowStackW) -> Result<Tensor<3>> {
589 let cfg = &self.config;
590 let xd = x.dims();
591 let (b_len, t_len) = (xd[0], xd[1]);
592 let num_cls = cfg.row_num_cls;
593 let out3 = self.row_interaction_3d(x, 0, w)?;
594 let e = cfg.embed_dim;
595 let sliced: Tensor<3> = out3.narrow(1, 0, num_cls);
596 let normed = rms(sliced, w.out_norm_w.clone(), cfg.norm_eps);
597 let icl_dim = num_cls * e;
598 Ok(normed.reshape([b_len, t_len, icl_dim]))
599 }
600
601 fn row_interaction_3d(
602 &self,
603 x: &Tensor<4>,
604 d_plus_cls: usize,
605 w: &BurnRowStackW,
606 ) -> Result<Tensor<3>> {
607 let cfg = &self.config;
608 let xd = x.dims();
609 let (b_len, t_len, hc, e) = (xd[0], xd[1], xd[2], xd[3]);
610 let mask = additive_key_mask(hc, d_plus_cls);
611 let rope = BurnTabfmRope {
612 freqs: w.rope_freqs.clone(),
613 };
614 let mut cur: Tensor<3> = x.clone().reshape([b_len * t_len, hc, e]);
615 for blk in &w.blocks {
616 cur = mab_forward(
617 &cur,
618 &cur,
619 &cur,
620 blk,
621 cfg.row_nhead,
622 cfg.norm_eps,
623 Some(mask.clone()),
624 Some(&rope),
625 )?;
626 }
627 Ok(cur)
628 }
629
630 fn icl_forward(
631 &self,
632 reps: &Tensor<3>,
633 y_batch: &[Vec<f32>],
634 train_size: usize,
635 ) -> Result<Tensor<3>> {
636 let cfg = &self.config;
637 let rd = reps.dims();
638 let (b_len, t_len) = (rd[0], rd[1]);
639 let y_enc = compute_y_encoder(y_batch, &self.icl.y_encoder, cfg.max_classes)?;
640 let tm: Vec<f32> = (0..t_len)
641 .map(|t| if t < train_size { 1.0 } else { 0.0 })
642 .collect();
643 let dev = device();
644 let tm_t = Tensor::<2>::from_data(TensorData::new(tm, [t_len, 1]), &dev);
645 let r = reps.clone() + y_enc * tm_t.unsqueeze_dim::<3>(0);
646 let mask = additive_key_mask(t_len, train_size);
647 let mut cur = r;
648 for blk in &self.icl.blocks {
649 cur = mab_forward(
650 &cur,
651 &cur,
652 &cur,
653 blk,
654 cfg.icl_nhead,
655 cfg.norm_eps,
656 Some(mask.clone()),
657 None,
658 )?;
659 }
660 let normed = rms(cur, self.icl.out_norm_w.clone(), cfg.norm_eps);
661 Ok(mlp_forward(&normed, &self.icl.decoder))
662 }
663}
664
665fn compute_y_embed(y_batch: &[Vec<f32>], w: &BurnYEmbedW, max_classes: usize) -> Result<Tensor<3>> {
666 let dev = device();
667 let b_len = y_batch.len();
668 let t = y_batch[0].len();
669 match w {
670 BurnYEmbedW::Embedding(emb_w) => {
671 let e = emb_w.dims()[1];
672 let emb: Vec<f32> = to_host(emb_w)?;
673 let mut out = vec![0f32; b_len * t * e];
674 for (bb, y) in y_batch.iter().enumerate() {
675 for i in 0..t {
676 let yi = y[i] as i64;
677 if yi >= 0 && (yi as usize) < max_classes {
678 let cls = yi as usize;
679 let dst = (bb * t + i) * e;
680 out[dst..dst + e].copy_from_slice(&emb[cls * e..(cls + 1) * e]);
681 }
682 }
683 }
684 Ok(Tensor::<3>::from_data(
685 TensorData::new(out, [b_len, t, e]),
686 &dev,
687 ))
688 }
689 BurnYEmbedW::Mlp(mlp) => {
690 let flat: Vec<f32> = y_batch.iter().flatten().copied().collect();
691 let y_col = Tensor::<2>::from_data(TensorData::new(flat, [b_len * t, 1]), &dev);
692 let out = mlp_forward(&y_col, mlp);
693 let e = out.dims()[1];
694 Ok(out.reshape([b_len, t, e]))
695 }
696 }
697}
698
699fn compute_y_encoder(
700 y_batch: &[Vec<f32>],
701 w: &BurnYEncoderW,
702 max_classes: usize,
703) -> Result<Tensor<3>> {
704 let dev = device();
705 let b_len = y_batch.len();
706 let t = y_batch[0].len();
707 match w {
708 BurnYEncoderW::OneHot { proj_w, proj_b } => {
709 let icl_dim = proj_w.dims()[0];
710 let w_flat: Vec<f32> = to_host(proj_w)?;
711 let bias: Vec<f32> = to_host_1(proj_b)?;
712 let mut out = vec![0f32; b_len * t * icl_dim];
713 for (bb, y) in y_batch.iter().enumerate() {
714 for i in 0..t {
715 let yi = y[i] as i64;
716 let valid = yi >= 0 && (yi as usize) < max_classes;
717 let dst = (bb * t + i) * icl_dim;
718 for e in 0..icl_dim {
719 let w_contrib = if valid {
720 w_flat[e * max_classes + yi as usize]
721 } else {
722 0.0
723 };
724 out[dst + e] = bias[e] + w_contrib;
725 }
726 }
727 }
728 Ok(Tensor::<3>::from_data(
729 TensorData::new(out, [b_len, t, icl_dim]),
730 &dev,
731 ))
732 }
733 BurnYEncoderW::Mlp(mlp) => {
734 let flat: Vec<f32> = y_batch.iter().flatten().copied().collect();
735 let y_col = Tensor::<2>::from_data(TensorData::new(flat, [b_len * t, 1]), &dev);
736 let out = mlp_forward(&y_col, mlp);
737 let icl_dim = out.dims()[1];
738 Ok(out.reshape([b_len, t, icl_dim]))
739 }
740 }
741}
742
743fn mlp_forward<const D: usize>(x: &Tensor<D>, w: &BurnMlpW) -> Tensor<D> {
744 let n = w.layers.len();
745 let mut h = x.clone();
746 for (i, (lw, lb)) in w.layers.iter().enumerate() {
747 h = linear(h, lw.clone(), Some(lb.clone()));
748 if i < n - 1 {
749 h = gelu_tanh(h);
750 }
751 }
752 h
753}
754
755fn additive_key_mask(seq_len: usize, valid_len: usize) -> Tensor<4> {
756 let dev = device();
757 let data: Vec<f32> = (0..seq_len)
758 .map(|i| if i < valid_len { 0.0 } else { -1e9 })
759 .collect();
760 Tensor::<1>::from_data(TensorData::new(data, [seq_len]), &dev).reshape([1, 1, 1, seq_len])
761}
762
763fn mab_forward(
764 q_raw: &Tensor<3>,
765 k_raw: &Tensor<3>,
766 v_raw: &Tensor<3>,
767 w: &BurnMabW,
768 nhead: usize,
769 eps: f64,
770 mask: Option<Tensor<4>>,
771 rope: Option<&BurnTabfmRope>,
772) -> Result<Tensor<3>> {
773 let qn = rms(q_raw.clone(), w.pre_attn_norm.clone(), eps);
774 let kn = rms(k_raw.clone(), w.pre_attn_norm.clone(), eps);
775 let vn = rms(v_raw.clone(), w.pre_attn_norm.clone(), eps);
776 let attn_out = mha_core(&qn, &kn, &vn, w, nhead, eps, mask, rope)?;
777 let a = rms(attn_out, w.post_attn_norm.clone(), eps);
778 let x = q_raw.clone() + a;
779 let xn = rms(x.clone(), w.pre_ff_norm.clone(), eps);
780 let gate = silu(linear(
781 xn.clone(),
782 w.ffn_gate_w.clone(),
783 Some(w.ffn_gate_b.clone()),
784 ));
785 let up = linear(xn, w.ffn_up_w.clone(), Some(w.ffn_up_b.clone()));
786 let ff = linear(gate * up, w.ffn_down_w.clone(), Some(w.ffn_down_b.clone()));
787 let ff = rms(ff, w.post_ff_norm.clone(), eps);
788 Ok(x + ff)
789}
790
791fn mha_core(
792 qn: &Tensor<3>,
793 kn: &Tensor<3>,
794 vn: &Tensor<3>,
795 w: &BurnMabW,
796 nhead: usize,
797 eps: f64,
798 mask: Option<Tensor<4>>,
799 rope: Option<&BurnTabfmRope>,
800) -> Result<Tensor<3>> {
801 let qd = qn.dims();
802 let (n, sq, e) = (qd[0], qd[1], qd[2]);
803 let sk = kn.dims()[1];
804 let hd = e / nhead;
805 let q = linear(qn.clone(), w.q_w.clone(), Some(w.q_b.clone())).reshape([n, sq, nhead, hd]);
806 let k = linear(kn.clone(), w.k_w.clone(), Some(w.k_b.clone())).reshape([n, sk, nhead, hd]);
807 let v = linear(vn.clone(), w.v_w.clone(), Some(w.v_b.clone())).reshape([n, sk, nhead, hd]);
808 let (q, k) = match rope {
809 Some(r) => (r.rotate(q)?, r.rotate(k)?),
810 None => (q, k),
811 };
812 let q = rms(q, w.q_norm.clone(), eps);
813 let k = rms(k, w.k_norm.clone(), eps);
814 let scale =
815 softplus(w.per_dim_scale.clone()).mul_scalar(std::f32::consts::LOG2_E / (hd as f32).sqrt());
816 let qd = q.dims();
817 let q = (q.reshape([qd[0] * qd[1] * qd[2], hd]) * scale.unsqueeze_dim::<2>(0)).reshape(qd);
818 let q: Tensor<4> = q.permute([0, 2, 1, 3]);
819 let k: Tensor<4> = k.permute([0, 2, 1, 3]);
820 let v: Tensor<4> = v.permute([0, 2, 1, 3]);
821 let mut scores = q.matmul(k.transpose());
822 if let Some(m) = mask {
823 scores = scores + m;
824 }
825 let probs = softmax_last(scores);
826 let out = probs.matmul(v);
827 let out: Tensor<4> = out.permute([0, 2, 1, 3]);
828 let out: Tensor<3> = out.reshape([n, sq, e]);
829 Ok(linear(out, w.o_w.clone(), Some(w.o_b.clone())))
830}
831
832#[cfg(test)]
833mod tests {
834 use super::*;
835 use crate::config::TabFMConfig;
836 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
837
838 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
839 let mut x = seed;
840 (0..n)
841 .map(|_| {
842 x = x
843 .wrapping_mul(6364136223846793005)
844 .wrapping_add(1442695040888963407);
845 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
846 })
847 .collect()
848 }
849
850 fn tiny_config_json() -> String {
851 serde_json::json!({
852 "embed_dim": 8,
853 "max_classes": 2,
854 "col_num_blocks": 1,
855 "col_nhead": 2,
856 "col_num_inds": 2,
857 "row_num_blocks": 1,
858 "row_nhead": 2,
859 "row_num_cls": 2,
860 "icl_num_blocks": 1,
861 "icl_nhead": 2,
862 "ff_factor": 2,
863 "feature_group_size": 2,
864 "is_classifier": true,
865 "num_freq": 2
866 })
867 .to_string()
868 }
869
870 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
871 *seed += 1;
872 let data: Vec<f32> = pseudo(rows * cols, *seed)
873 .iter()
874 .map(|v| v * scale)
875 .collect();
876 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
877 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
878 }
879
880 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
881 *seed += 1;
882 let bytes: Vec<u8> = pseudo(n, *seed)
883 .iter()
884 .flat_map(|v| v.to_le_bytes())
885 .collect();
886 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
887 }
888
889 fn put_mab(w: &mut GGUFWriter, prefix: &str, e: usize, ff: usize, seed: &mut u64, s: f32) {
890 for t in ["attn_q", "attn_k", "attn_v", "attn_o"] {
891 put2(w, &format!("{prefix}.{t}.weight"), e, e, seed, s);
892 put1(w, &format!("{prefix}.{t}.bias"), e, seed);
893 }
894 put1(w, &format!("{prefix}.q_norm.weight"), e / 2, seed);
895 put1(w, &format!("{prefix}.k_norm.weight"), e / 2, seed);
896 put1(w, &format!("{prefix}.per_dim_scale"), e / 2, seed);
897 for t in [
898 "pre_attn_norm",
899 "post_attn_norm",
900 "pre_ff_norm",
901 "post_ff_norm",
902 ] {
903 put1(w, &format!("{prefix}.{t}.weight"), e, seed);
904 }
905 put2(w, &format!("{prefix}.ffn_up.weight"), ff, e, seed, s);
906 put1(w, &format!("{prefix}.ffn_up.bias"), ff, seed);
907 put2(w, &format!("{prefix}.ffn_gate.weight"), ff, e, seed, s);
908 put1(w, &format!("{prefix}.ffn_gate.bias"), ff, seed);
909 put2(w, &format!("{prefix}.ffn_down.weight"), e, ff, seed, s);
910 put1(w, &format!("{prefix}.ffn_down.bias"), e, seed);
911 }
912
913 #[test]
914 fn burn_single_table_golden() {
915 let dir = std::env::temp_dir().join(format!("tabfm-burn-{}", std::process::id()));
916 let _ = std::fs::remove_dir_all(&dir);
917 std::fs::create_dir_all(&dir).unwrap();
918 let path = dir.join("model.gguf");
919 let mut w = GGUFWriter::new();
920 w.add_metadata(
921 "general.architecture",
922 GGUFMetaValue::String("tabfm".into()),
923 );
924 let mut seed = 2500u64;
925 let s = 0.05;
926 put2(&mut w, "cell.fourier_freq", 2, 2, &mut seed, 1.0);
927 put2(&mut w, "cell.fourier_freq_cat", 2, 2, &mut seed, 1.0);
928 put2(&mut w, "cell.in_linear.weight", 8, 4, &mut seed, s);
929 put1(&mut w, "cell.in_linear.bias", 8, &mut seed);
930 put2(&mut w, "cell.in_linear_cat.weight", 8, 4, &mut seed, s);
931 put1(&mut w, "cell.in_linear_cat.bias", 8, &mut seed);
932 put2(&mut w, "cell.y_embed.weight", 2, 8, &mut seed, s);
933 for stack in ["colenc1", "colenc2"] {
934 put2(
935 &mut w,
936 &format!("{stack}.blk.0.ind_vectors"),
937 2,
938 8,
939 &mut seed,
940 s,
941 );
942 put_mab(&mut w, &format!("{stack}.blk.0.mab1"), 8, 16, &mut seed, s);
943 put_mab(&mut w, &format!("{stack}.blk.0.mab2"), 8, 16, &mut seed, s);
944 put2(&mut w, &format!("{stack}.out_w.weight"), 8, 8, &mut seed, s);
945 put1(&mut w, &format!("{stack}.out_w.bias"), 8, &mut seed);
946 put1(&mut w, &format!("{stack}.out_norm.weight"), 8, &mut seed);
947 }
948 for stack in ["rowenc1", "rowenc2"] {
949 put1(&mut w, &format!("{stack}.rope_freqs"), 2, &mut seed);
950 put_mab(&mut w, &format!("{stack}.blk.0"), 8, 16, &mut seed, s);
951 put1(&mut w, &format!("{stack}.out_norm.weight"), 8, &mut seed);
952 }
953 put2(&mut w, "cls_tokens", 2, 8, &mut seed, s);
954 put_mab(&mut w, "icl.blk.0", 16, 32, &mut seed, s);
955 put2(
956 &mut w,
957 "icl.y_encoder.projection.weight",
958 16,
959 2,
960 &mut seed,
961 s,
962 );
963 put1(&mut w, "icl.y_encoder.projection.bias", 16, &mut seed);
964 put2(&mut w, "icl.decoder.mlp.0.weight", 32, 16, &mut seed, s);
965 put1(&mut w, "icl.decoder.mlp.0.bias", 32, &mut seed);
966 put2(&mut w, "icl.decoder.mlp.1.weight", 2, 32, &mut seed, s);
967 put1(&mut w, "icl.decoder.mlp.1.bias", 2, &mut seed);
968 put1(&mut w, "icl.out_norm.weight", 16, &mut seed);
969 let f = std::fs::File::create(&path).unwrap();
970 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
971
972 let tc = TabFMConfig::from_json(&tiny_config_json()).unwrap();
973 let burn = BurnTabFMModel::load(&path, &tc).unwrap();
974 let x = vec![
975 vec![0.1, 1.2],
976 vec![0.9, -0.3],
977 vec![-1.1, 0.4],
978 vec![1.5, 1.1],
979 vec![0.2, 0.9],
980 vec![-0.8, 0.1],
981 ];
982 let y = vec![0.0, 1.0, 1.0, 0.0, 0.0, 0.0];
983 let t0 = std::time::Instant::now();
984 let a = burn.predict(&x, &y, 4, None, None).unwrap();
985 let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
986 let b = burn.predict(&x, &y, 4, None, None).unwrap();
987 assert_eq!(a.len(), 6);
988 assert!(a.iter().all(|r| r.iter().all(|v| v.is_finite())));
989 assert_eq!(a, b, "predict must be deterministic");
990 println!("tabfm synthetic: burn {burn_ms:.1}ms out={a:?}");
991 let golden = [
993 [-0.6438628, 0.015820652],
994 [-0.6436042, 0.016053759],
995 [-0.6436042, 0.016053751],
996 [-0.6438628, 0.015820652],
997 [-0.6616832, -0.001892142],
998 [-0.6616832, -0.0018921569],
999 ];
1000 for (ra, re) in a.iter().zip(golden.iter()) {
1001 for (u, v) in ra.iter().zip(re.iter()) {
1002 assert!((u - v).abs() < 1e-3, "golden drift: {u} vs {v}");
1003 }
1004 }
1005 let _ = std::fs::remove_dir_all(&dir);
1006 }
1007}