Skip to main content

zsfm_timesfm/
tensor_map.rs

1/// Map a HuggingFace TimesFM tensor name to its GGUF blk.N.* equivalent.
2///
3/// HF parameter naming comes from PyTorch nn.Module traversal of
4/// `TimesFM_2p5_200M_torch_module`. Returns `None` for any name that
5/// should not be included in the GGUF (currently none are skipped).
6pub fn map_tensor_name(hf_name: &str) -> Option<String> {
7    // Tokenizer (ResidualBlock, with bias)
8    match hf_name {
9        "tokenizer.hidden_layer.weight"   => return Some("tokenizer.hidden.weight".into()),
10        "tokenizer.hidden_layer.bias"     => return Some("tokenizer.hidden.bias".into()),
11        "tokenizer.output_layer.weight"   => return Some("tokenizer.output.weight".into()),
12        "tokenizer.output_layer.bias"     => return Some("tokenizer.output.bias".into()),
13        "tokenizer.residual_layer.weight" => return Some("tokenizer.skip.weight".into()),
14        "tokenizer.residual_layer.bias"   => return Some("tokenizer.skip.bias".into()),
15        // Output projection — point head (ResidualBlock, no bias)
16        "output_projection_point.hidden_layer.weight"   => return Some("out_point.hidden.weight".into()),
17        "output_projection_point.output_layer.weight"   => return Some("out_point.output.weight".into()),
18        "output_projection_point.residual_layer.weight" => return Some("out_point.skip.weight".into()),
19        // Output projection — quantile head (ResidualBlock, no bias; stored but not used in basic infer)
20        "output_projection_quantiles.hidden_layer.weight"   => return Some("out_quantile.hidden.weight".into()),
21        "output_projection_quantiles.output_layer.weight"   => return Some("out_quantile.output.weight".into()),
22        "output_projection_quantiles.residual_layer.weight" => return Some("out_quantile.skip.weight".into()),
23        _ => {}
24    }
25
26    // Transformer blocks: stacked_xf.{N}.*
27    let rest = hf_name.strip_prefix("stacked_xf.")?;
28    let (block_str, rest) = rest.split_once('.')?;
29    let block: u32 = block_str.parse().ok()?;
30
31    let gguf_suffix = match rest {
32        "pre_attn_ln.scale"                       => "pre_attn_norm.weight",
33        "post_attn_ln.scale"                      => "post_attn_norm.weight",
34        "attn.qkv_proj.weight"                    => "attn_qkv.weight",
35        "attn.out.weight"                         => "attn_out.weight",
36        "attn.query_ln.scale"                     => "attn_q_norm.weight",
37        "attn.key_ln.scale"                       => "attn_k_norm.weight",
38        "attn.per_dim_scale.per_dim_scale"        => "attn_q_scale.weight",
39        "pre_ff_ln.scale"                         => "pre_ff_norm.weight",
40        "post_ff_ln.scale"                        => "post_ff_norm.weight",
41        "ff0.weight"                              => "ffn_up.weight",
42        "ff1.weight"                              => "ffn_down.weight",
43        _ => return None,
44    };
45
46    Some(format!("blk.{block}.{gguf_suffix}"))
47}