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 *default-ish* Python TabICL config:
//! `col_target_aware=true`, `col_feature_group="same"`, exercising the
//! target-aware embedding + circular-permutation feature grouping paths.

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_default_config_forward_matches_rust() {
    let ckpt_path = Path::new("/tmp/tabicl_parity_default/ckpt.json");
    let input_path = Path::new("/tmp/tabicl_parity_default/input.bin");
    let output_path = Path::new("/tmp/tabicl_parity_default/output.bin");

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

    let mut cfg = TabICLConfig::default();
    cfg.max_classes = 3;
    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::Same;
    cfg.col_feature_group_size = 3;
    cfg.col_target_aware = true;
    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 default-config checkpoint");

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

    // X shape: (1, 5, 4); y_train: (1, 3).
    let x = ndarray::Array3::from_shape_vec((1, 5, 4), x_flat).unwrap();
    let y_train = ndarray::Array2::from_shape_vec((1, 3), vec![0_usize, 1, 2]).unwrap();

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

    // Python reference: shape (1, n_test, max_classes) — n_test = 2.
    let n_train = 3;
    let n_test = y_ref.len() / cfg.max_classes;
    eprintln!("Python ref: {:?}", y_ref);
    let mut rust_test_out = Vec::new();
    for i in 0..n_test {
        for c in 0..cfg.max_classes {
            rust_test_out.push(out[(0, n_train + i, c)]);
        }
    }
    eprintln!("Rust out (test rows): {:?}", rust_test_out);

    let mut max_diff = 0.0_f32;
    for (rust_val, py_val) in rust_test_out.iter().zip(y_ref.iter()) {
        let diff = (rust_val - py_val).abs();
        if diff > max_diff {
            max_diff = diff;
        }
    }
    eprintln!("Max diff: {max_diff}");
    assert!(
        max_diff < 1e-5,
        "default-config parity exceeded tolerance: max_diff={max_diff}"
    );
}