zsfm_tirex/
tensor_map.rs

1/// Map a TiRex state_dict tensor name (from PyTorch Lightning .ckpt) to GGUF naming.
2///
3/// The checkpoint stores block weights under "block_stack.blocks.N.*" and the
4/// non-block embeddings / output norm without a "block_stack." prefix.
5pub fn map_tensor_name(name: &str) -> Option<String> {
6    // Non-block tensors
7    match name {
8        "block_stack.out_norm.weight" => return Some("out_norm".into()),
9        _ => {}
10    }
11
12    // input/output patch embedding (no block_stack prefix)
13    if let Some(rest) = name.strip_prefix("input_patch_embedding.") {
14        let suffix = match rest {
15            "hidden_layer.weight"   => "in_emb.hidden.weight",
16            "hidden_layer.bias"     => "in_emb.hidden.bias",
17            "output_layer.weight"   => "in_emb.output.weight",
18            "output_layer.bias"     => "in_emb.output.bias",
19            "residual_layer.weight" => "in_emb.residual.weight",
20            "residual_layer.bias"   => "in_emb.residual.bias",
21            _ => return None,
22        };
23        return Some(suffix.into());
24    }
25    if let Some(rest) = name.strip_prefix("output_patch_embedding.") {
26        let suffix = match rest {
27            "hidden_layer.weight"   => "out_emb.hidden.weight",
28            "hidden_layer.bias"     => "out_emb.hidden.bias",
29            "output_layer.weight"   => "out_emb.output.weight",
30            "output_layer.bias"     => "out_emb.output.bias",
31            "residual_layer.weight" => "out_emb.residual.weight",
32            "residual_layer.bias"   => "out_emb.residual.bias",
33            _ => return None,
34        };
35        return Some(suffix.into());
36    }
37
38    // block_stack.blocks.N.*
39    if let Some(rest) = name.strip_prefix("block_stack.blocks.") {
40        let (n_str, rest) = rest.split_once('.')?;
41        let n: u32 = n_str.parse().ok()?;
42        let gguf_suffix = match rest {
43            "norm_slstm.weight"                   => "norm_slstm",
44            "slstm_layer.fgate.weight"             => "fgate.weight",
45            "slstm_layer.igate.weight"             => "igate.weight",
46            "slstm_layer.zgate.weight"             => "zgate.weight",
47            "slstm_layer.ogate.weight"             => "ogate.weight",
48            "slstm_layer.slstm_cell._recurrent_kernel_" => "slstm_kernel",
49            "slstm_layer.slstm_cell._bias_"        => "slstm_bias",
50            "slstm_layer.group_norm.weight"        => "group_norm",
51            "norm_ffn.weight"                      => "norm_ffn",
52            "ffn.proj_up_gate.weight"              => "ffn_gate.weight",
53            "ffn.proj_up.weight"                   => "ffn_up.weight",
54            "ffn.proj_down.weight"                 => "ffn_down.weight",
55            _ => return None,
56        };
57        return Some(format!("blk.{n}.{gguf_suffix}"));
58    }
59
60    None
61}
62
63/// Whether this tensor needs bias permutation ([NH, NG, DH] → [NG, NH, DH]).
64pub fn needs_bias_permute(gguf_name: &str) -> bool {
65    gguf_name.ends_with(".slstm_bias")
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn non_block() {
74        assert_eq!(
75            map_tensor_name("block_stack.out_norm.weight"),
76            Some("out_norm".into())
77        );
78        assert_eq!(
79            map_tensor_name("input_patch_embedding.hidden_layer.weight"),
80            Some("in_emb.hidden.weight".into())
81        );
82        assert_eq!(
83            map_tensor_name("output_patch_embedding.residual_layer.bias"),
84            Some("out_emb.residual.bias".into())
85        );
86    }
87
88    #[test]
89    fn block() {
90        assert_eq!(
91            map_tensor_name("block_stack.blocks.0.norm_slstm.weight"),
92            Some("blk.0.norm_slstm".into())
93        );
94        assert_eq!(
95            map_tensor_name("block_stack.blocks.11.slstm_layer.slstm_cell._bias_"),
96            Some("blk.11.slstm_bias".into())
97        );
98        assert_eq!(
99            map_tensor_name("block_stack.blocks.3.ffn.proj_down.weight"),
100            Some("blk.3.ffn_down.weight".into())
101        );
102    }
103}
zsfm_tirex/tensor_map.rs

zsfm_tirex/
tensor_map.rs
