tabicl-model 2.1.1

TabICL transformer model — column embedding, row interaction, ICL learning, KV cache.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

//! Cross-stack parity for the regression head (`max_classes = 0`, the
//! quantile-prediction path). Verifies the Rust forward output matches
//! Python `tabicl._model.TabICL._train_forward` on a 5-quantile head.

use std::path::Path;
use tabicl_model::TabICL;
use tabicl_model::tabicl::{ColFeatureGroup, TabICLConfig};

fn read_f32_blob(path: &Path) -> Option<Vec<f32>> {
    let bytes = std::fs::read(path).ok()?;
    let mut out = Vec::with_capacity(bytes.len() / 4);
    for chunk in bytes.chunks_exact(4) {
        out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
    }
    Some(out)
}

#[test]
fn python_regression_forward_matches_rust() {
    let ckpt_path = Path::new("/tmp/tabicl_parity_reg/ckpt.json");
    let input_path = Path::new("/tmp/tabicl_parity_reg/input.bin");
    let y_train_path = Path::new("/tmp/tabicl_parity_reg/y_train.bin");
    let output_path = Path::new("/tmp/tabicl_parity_reg/output.bin");

    if !ckpt_path.exists() {
        eprintln!(
            "regression parity: fixture not found at {:?}, skipping",
            ckpt_path
        );
        return;
    }

    let mut cfg = TabICLConfig::default();
    cfg.max_classes = 0; // regression mode
    cfg.num_quantiles = 5;
    cfg.embed_dim = 8;
    cfg.col_num_blocks = 1;
    cfg.col_nhead = 2;
    cfg.col_num_inds = 4;
    cfg.col_affine = false;
    cfg.col_feature_group = ColFeatureGroup::None;
    cfg.col_target_aware = false;
    cfg.col_ssmax = "none".into();
    cfg.row_num_blocks = 1;
    cfg.row_nhead = 2;
    cfg.row_num_cls = 2;
    cfg.row_rope_base = 100_000.0;
    cfg.row_rope_interleaved = false;
    cfg.icl_num_blocks = 1;
    cfg.icl_nhead = 2;
    cfg.icl_ssmax = "none".into();
    cfg.ff_factor = 2;
    cfg.dropout = 0.0;
    cfg.norm_first = true;
    cfg.bias_free_ln = false;

    let mut model = TabICL::new(cfg.clone());
    model
        .load_from_file(ckpt_path)
        .expect("load Python regression checkpoint");

    let x_flat = read_f32_blob(input_path).expect("read input");
    let y_train_flat = read_f32_blob(y_train_path).expect("read y_train");
    let y_ref = read_f32_blob(output_path).expect("read output");

    let x = ndarray::Array3::from_shape_vec((1, 4, 3), x_flat).unwrap();
    let y_train_arr = ndarray::Array2::from_shape_vec((1, 3), y_train_flat).unwrap();

    let out = model
        .forward(x.view(), None, Some(y_train_arr.view()))
        .expect("Rust forward");

    // Python ref shape: (1, 1, 5) — the test row's quantile predictions.
    let n_train = 3;
    let n_test = y_ref.len() / cfg.num_quantiles;
    eprintln!("Python ref (quantiles): {:?}", y_ref);
    let mut rust_test: Vec<f32> = Vec::new();
    for i in 0..n_test {
        for q in 0..cfg.num_quantiles {
            rust_test.push(out[(0, n_train + i, q)]);
        }
    }
    eprintln!("Rust quantiles (test): {:?}", rust_test);

    let mut max_diff = 0.0_f32;
    for (r, p) in rust_test.iter().zip(y_ref.iter()) {
        let d = (r - p).abs();
        if d > max_diff {
            max_diff = d;
        }
    }
    eprintln!("Max diff: {max_diff}");
    assert!(
        max_diff < 1e-5,
        "regression parity exceeded tolerance: max_diff={max_diff}"
    );
}