ariacompute-inference 1.17.1

//! Tensor name aliases across Aria families.
//!
//! Tiny/`--tiny` uses GGUF-style `blk.*`. Real `model` quantize keeps HuggingFace
//! state_dict names, which vary by family:
//! - Qwen / Gemma-3 text / LFM / Nanbeige / …: `model.layers.*`
//! - Gemma-4 / Gemma-3n / LFM-VL / …: `model.language_model.layers.*`
//! - OpenVLA-style: `language_model.model.layers.*`
//! - Gemma-2+ FFN pre-norm: `pre_feedforward_layernorm` (in addition to
//!   LLaMA/Qwen `post_attention_layernorm`)

/// HF-style decoder layer prefixes (without trailing layer index).
const HF_LAYER_PREFIXES: &[&str] = &[
    "model.layers",
    "model.language_model.layers",
    "language_model.model.layers",
    "model.model.layers",
];

fn with_layer_suffix(layer: usize, suffix: &str) -> Vec<String> {
    HF_LAYER_PREFIXES
        .iter()
        .map(|p| format!("{p}.{layer}.{suffix}"))
        .collect()
}

/// Candidates for a logical weight (first hit wins).
pub fn emb_names() -> Vec<&'static str> {
    vec![
        "token_embd.weight",
        // Gemma-4 / VL nested path before plain `model.embed_tokens` (Qwen).
        "model.language_model.embed_tokens.weight",
        "language_model.model.embed_tokens.weight",
        "model.embed_tokens.weight",
        "model.model.embed_tokens.weight",
    ]
}

pub fn output_norm_names() -> Vec<&'static str> {
    vec![
        "output_norm.weight",
        "model.norm.weight",
        "model.language_model.norm.weight",
        "language_model.model.norm.weight",
        "model.model.norm.weight",
    ]
}

/// LM head; last entries are weight-tied embedding fallbacks.
pub fn output_names() -> Vec<&'static str> {
    vec![
        "output.weight",
        "lm_head.weight",
        "model.language_model.embed_tokens.weight",
        "language_model.model.embed_tokens.weight",
        "model.embed_tokens.weight",
        "model.model.embed_tokens.weight",
    ]
}

pub fn attn_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_norm.weight")];
    names.extend(with_layer_suffix(layer, "input_layernorm.weight"));
    // LFM2/2.5: operator_norm before attn or short-conv.
    names.extend(with_layer_suffix(layer, "operator_norm.weight"));
    names
}

pub fn ffn_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_norm.weight")];
    // Gemma-2 / Gemma-3 / Gemma-4: dedicated pre-FFN norm first.
    names.extend(pre_feedforward_norm_names(layer));
    // LLaMA / Qwen: post-attention = pre-FFN.
    names.extend(with_layer_suffix(layer, "post_attention_layernorm.weight"));
    // LFM2 dense / conv layers.
    names.extend(with_layer_suffix(layer, "ffn_norm.weight"));
    names
}

/// Gemma-2+ dedicated pre-FFN norm (presence implies the 4-norm residual graph).
pub fn pre_feedforward_norm_names(layer: usize) -> Vec<String> {
    with_layer_suffix(layer, "pre_feedforward_layernorm.weight")
}

/// Gemma-2+ residual post-attn norm (not the LLaMA FFN norm alias).
pub fn attn_post_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_post_norm.weight")];
    names.extend(with_layer_suffix(layer, "post_attention_layernorm.weight"));
    names
}

/// Gemma-2+ residual post-FFN norm.
pub fn ffn_post_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_post_norm.weight")];
    names.extend(with_layer_suffix(
        layer,
        "post_feedforward_layernorm.weight",
    ));
    names
}

pub fn embed_per_layer_names() -> Vec<&'static str> {
    vec![
        "model.language_model.embed_tokens_per_layer.weight",
        "language_model.model.embed_tokens_per_layer.weight",
        "model.embed_tokens_per_layer.weight",
    ]
}

pub fn per_layer_model_projection_names() -> Vec<&'static str> {
    vec![
        "model.per_layer_model_projection.weight",
        "model.language_model.per_layer_model_projection.weight",
        "language_model.model.per_layer_model_projection.weight",
    ]
}

pub fn per_layer_projection_norm_names() -> Vec<&'static str> {
    vec![
        "model.per_layer_projection_norm.weight",
        "model.language_model.per_layer_projection_norm.weight",
        "language_model.model.per_layer_projection_norm.weight",
    ]
}

pub fn layer_ple_gate_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ple_gate.weight")];
    names.extend(with_layer_suffix(layer, "per_layer_input_gate.weight"));
    names
}

pub fn layer_ple_proj_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ple_proj.weight")];
    names.extend(with_layer_suffix(layer, "per_layer_projection.weight"));
    names
}

pub fn layer_ple_post_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ple_post_norm.weight")];
    names.extend(with_layer_suffix(layer, "post_per_layer_input_norm.weight"));
    names
}

/// HF Gemma-4 `layers.{i}.layer_scalar` (converted from JAX `skip_scale`).
/// Not a `.weight` tensor — a 0-d / `[1]` residual multiplier after the layer.
pub fn layer_scalar_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.layer_scalar")];
    names.extend(with_layer_suffix(layer, "layer_scalar"));
    names
}

/// Gemma-3n AltUp / Laurel (HF `Gemma3nTextDecoderLayer`).
pub fn layer_altup_router_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.altup_router.weight")];
    names.extend(with_layer_suffix(layer, "altup.modality_router.weight"));
    names
}

pub fn layer_altup_router_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.altup_router_norm.weight")];
    names.extend(with_layer_suffix(layer, "altup.router_norm.weight"));
    names
}

pub fn layer_altup_prediction_coef_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.altup_pred.weight")];
    names.extend(with_layer_suffix(layer, "altup.prediction_coefs.weight"));
    names
}

pub fn layer_altup_correction_coef_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.altup_corr.weight")];
    names.extend(with_layer_suffix(layer, "altup.correction_coefs.weight"));
    names
}

pub fn layer_altup_correct_scale_names(layer: usize) -> Vec<String> {
    let mut names = vec![
        format!("blk.{layer}.altup_correct_scale"),
        format!("blk.{layer}.altup_correct_scale.weight"),
    ];
    names.extend(with_layer_suffix(layer, "altup.correct_output_scale"));
    names.extend(with_layer_suffix(
        layer,
        "altup.correct_output_scale.weight",
    ));
    names
}

pub fn layer_laurel_left_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.laurel_left.weight")];
    names.extend(with_layer_suffix(layer, "laurel.linear_left.weight"));
    names
}

pub fn layer_laurel_right_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.laurel_right.weight")];
    names.extend(with_layer_suffix(layer, "laurel.linear_right.weight"));
    names
}

pub fn layer_laurel_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.laurel_norm.weight")];
    names.extend(with_layer_suffix(layer, "laurel.post_laurel_norm.weight"));
    names
}

fn model_level_names(suffix: &str) -> Vec<String> {
    [
        format!("model.language_model.{suffix}"),
        format!("language_model.model.{suffix}"),
        format!("model.{suffix}"),
    ]
    .into_iter()
    .collect()
}

pub fn altup_projection_names(index: usize) -> Vec<String> {
    let mut names = model_level_names(&format!("altup_projections.{index}.weight"));
    names.insert(0, format!("altup_proj.{index}.weight"));
    names
}

pub fn altup_unembed_names(index: usize) -> Vec<String> {
    let mut names = model_level_names(&format!("altup_unembed_projections.{index}.weight"));
    names.insert(0, format!("altup_unembed.{index}.weight"));
    names
}

pub fn attn_q_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_q_norm.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.q_norm.weight"));
    names.extend(with_layer_suffix(layer, "self_attn.query_layernorm.weight"));
    names
}

pub fn attn_k_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_k_norm.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.k_norm.weight"));
    names.extend(with_layer_suffix(layer, "self_attn.key_layernorm.weight"));
    names
}

pub fn attn_v_norm_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_v_norm.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.v_norm.weight"));
    names.extend(with_layer_suffix(layer, "self_attn.value_layernorm.weight"));
    names
}

pub fn attn_q_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_q.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.q_proj.weight"));
    names
}

pub fn attn_k_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_k.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.k_proj.weight"));
    names
}

pub fn attn_v_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_v.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.v_proj.weight"));
    names
}

pub fn attn_o_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.attn_output.weight")];
    names.extend(with_layer_suffix(layer, "self_attn.o_proj.weight"));
    names
}

pub fn ffn_gate_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_gate.weight")];
    names.extend(with_layer_suffix(layer, "mlp.gate_proj.weight"));
    names.extend(with_layer_suffix(layer, "mlp.w1.weight"));
    names.extend(with_layer_suffix(layer, "feed_forward.w1.weight"));
    names
}

pub fn ffn_up_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_up.weight")];
    names.extend(with_layer_suffix(layer, "mlp.up_proj.weight"));
    names.extend(with_layer_suffix(layer, "mlp.w3.weight"));
    names.extend(with_layer_suffix(layer, "feed_forward.w3.weight"));
    names
}

pub fn ffn_down_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_down.weight")];
    names.extend(with_layer_suffix(layer, "mlp.down_proj.weight"));
    names.extend(with_layer_suffix(layer, "mlp.w2.weight"));
    names.extend(with_layer_suffix(layer, "feed_forward.w2.weight"));
    names
}

/// LFM2 short-conv `in_proj` (projects to 3×hidden: B, C, x).
pub fn conv_in_proj_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.conv_in.weight")];
    names.extend(with_layer_suffix(layer, "conv.in_proj.weight"));
    names
}

pub fn conv_out_proj_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.conv_out.weight")];
    names.extend(with_layer_suffix(layer, "conv.out_proj.weight"));
    names
}

/// Depthwise Conv1d weight; often stored as `[hidden, kernel]` after squeeze.
pub fn conv_kernel_names(layer: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.conv_kernel.weight")];
    names.extend(with_layer_suffix(layer, "conv.conv.weight"));
    names
}

/// MoE router / gate (logits over experts).
pub fn moe_router_names(layer: usize) -> Vec<String> {
    let mut names = vec![
        format!("blk.{layer}.ffn_gate_inp.weight"),
        format!("blk.{layer}.moe_gate.weight"),
    ];
    names.extend(with_layer_suffix(layer, "block_sparse_moe.gate.weight"));
    names.extend(with_layer_suffix(layer, "mlp.gate.weight"));
    names.extend(with_layer_suffix(layer, "feed_forward.gate.weight"));
    names
}

pub fn moe_expert_gate_names(layer: usize, expert: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_gate.{expert}.weight")];
    names.extend(with_layer_suffix(
        layer,
        &format!("block_sparse_moe.experts.{expert}.w1.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("block_sparse_moe.experts.{expert}.gate_proj.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("mlp.experts.{expert}.gate_proj.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("mlp.experts.{expert}.w1.weight"),
    ));
    names
}

pub fn moe_expert_up_names(layer: usize, expert: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_up.{expert}.weight")];
    names.extend(with_layer_suffix(
        layer,
        &format!("block_sparse_moe.experts.{expert}.w3.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("block_sparse_moe.experts.{expert}.up_proj.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("mlp.experts.{expert}.up_proj.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("mlp.experts.{expert}.w3.weight"),
    ));
    names
}

pub fn moe_expert_down_names(layer: usize, expert: usize) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.ffn_down.{expert}.weight")];
    names.extend(with_layer_suffix(
        layer,
        &format!("block_sparse_moe.experts.{expert}.w2.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("block_sparse_moe.experts.{expert}.down_proj.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("mlp.experts.{expert}.down_proj.weight"),
    ));
    names.extend(with_layer_suffix(
        layer,
        &format!("mlp.experts.{expert}.w2.weight"),
    ));
    names
}

/// Qwen3.5 / Bonsai Gated DeltaNet projections.
/// Qwen3-Next fused `in_proj_qkvz` / `in_proj_ba`; Qwen3.5 splits them into
/// `in_proj_qkv`+`in_proj_z` and `in_proj_b`+`in_proj_a`.
pub fn linear_in_proj_qkvz_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "in_proj_qkvz.weight")
}

pub fn linear_in_proj_qkv_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "in_proj_qkv.weight")
}

pub fn linear_in_proj_z_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "in_proj_z.weight")
}

pub fn linear_in_proj_ba_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "in_proj_ba.weight")
}

pub fn linear_in_proj_b_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "in_proj_b.weight")
}

pub fn linear_in_proj_a_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "in_proj_a.weight")
}

pub fn linear_conv1d_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "conv1d.weight")
}

pub fn linear_out_proj_names(layer: usize) -> Vec<String> {
    linear_attn_names(layer, "out_proj.weight")
}

/// Gated DeltaNet output RMSNormGated (`* w`, ones-init).
pub fn linear_out_norm_names(layer: usize) -> Vec<String> {
    let mut names = linear_attn_names(layer, "norm.weight");
    names.extend(linear_attn_names(layer, "out_norm.weight"));
    names
}

pub fn linear_a_log_names(layer: usize) -> Vec<String> {
    let mut names = linear_attn_names(layer, "A_log");
    names.extend(linear_attn_names(layer, "_fp32_params.A_log"));
    names
}

pub fn linear_dt_bias_names(layer: usize) -> Vec<String> {
    let mut names = linear_attn_names(layer, "dt_bias");
    names.extend(linear_attn_names(layer, "_fp32_params.dt_bias"));
    names
}

fn linear_attn_names(layer: usize, suffix: &str) -> Vec<String> {
    let mut names = vec![format!("blk.{layer}.linear_attn.{suffix}")];
    names.extend(with_layer_suffix(layer, &format!("linear_attn.{suffix}")));
    names.extend(with_layer_suffix(layer, &format!("self_attn.{suffix}")));
    names
}

/// Vision projector / merger (first hit wins at load).
pub fn vision_proj_names() -> Vec<&'static str> {
    vec![
        "mm_projector.weight",
        "vision_projector.weight",
        "multi_modal_projector.linear.weight",
        "multi_modal_projector.weight",
        "model.visual.merger.mlp.0.weight",
        "visual.merger.mlp.0.weight",
        "model.mm_projector.weight",
        "vision_tower.projector.weight",
    ]
}

/// VLA action / policy head.
pub fn action_head_names() -> Vec<&'static str> {
    vec![
        "action_head.weight",
        "action_out.weight",
        "model.action_head.weight",
        "model.action_out_proj.weight",
        "policy.action_out.weight",
    ]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn aliases_cover_hf_blk_and_language_model() {
        let a = attn_norm_names(0);
        assert!(a.iter().any(|s| s.contains("attn_norm")));
        assert!(a
            .iter()
            .any(|s| s == "model.layers.0.input_layernorm.weight"));
        assert!(a.iter().any(|s| s == "model.layers.0.operator_norm.weight"));
        assert!(a
            .iter()
            .any(|s| s == "model.language_model.layers.0.input_layernorm.weight"));
        assert!(a
            .iter()
            .any(|s| s == "language_model.model.layers.0.input_layernorm.weight"));

        let f = ffn_norm_names(0);
        assert!(f
            .iter()
            .any(|s| s == "model.layers.0.post_attention_layernorm.weight"));
        assert!(f
            .iter()
            .any(|s| s == "model.layers.0.pre_feedforward_layernorm.weight"));
        assert!(f
            .iter()
            .any(|s| s.contains("language_model") && s.contains("pre_feedforward_layernorm")));

        assert_eq!(
            emb_names()[1],
            "model.language_model.embed_tokens.weight",
            "Gemma-4 nested embed must be tried before Qwen model.embed_tokens"
        );
        assert!(emb_names().contains(&"model.embed_tokens.weight"));
        assert!(emb_names().contains(&"model.language_model.embed_tokens.weight"));
        assert!(output_names().contains(&"lm_head.weight"));
        assert!(output_norm_names().contains(&"model.language_model.norm.weight"));
        assert!(
            embed_per_layer_names().contains(&"model.language_model.embed_tokens_per_layer.weight")
        );
        assert!(attn_post_norm_names(0)
            .iter()
            .any(|s| s == "model.layers.0.post_attention_layernorm.weight"));
        assert!(ffn_post_norm_names(0)
            .iter()
            .any(|s| s.contains("post_feedforward_layernorm")));
        assert!(layer_ple_gate_names(0)
            .iter()
            .any(|s| s.contains("per_layer_input_gate.weight")));
        assert!(layer_scalar_names(0)
            .iter()
            .any(|s| s == "model.language_model.layers.0.layer_scalar"));
        assert!(layer_altup_router_names(0)
            .iter()
            .any(|s| s == "model.language_model.layers.0.altup.modality_router.weight"));
        assert!(layer_laurel_left_names(0)
            .iter()
            .any(|s| s == "model.language_model.layers.0.laurel.linear_left.weight"));
        assert!(altup_projection_names(0)
            .iter()
            .any(|s| s == "model.language_model.altup_projections.0.weight"));
        assert!(altup_unembed_names(1)
            .iter()
            .any(|s| s == "model.language_model.altup_unembed_projections.1.weight"));
        assert!(linear_in_proj_qkv_names(0)
            .iter()
            .any(|s| s == "model.layers.0.linear_attn.in_proj_qkv.weight"));
        assert!(linear_in_proj_z_names(0)
            .iter()
            .any(|s| s == "model.language_model.layers.0.linear_attn.in_proj_z.weight"));
        assert!(linear_in_proj_b_names(0)
            .iter()
            .any(|s| s == "model.layers.0.linear_attn.in_proj_b.weight"));
        assert!(linear_in_proj_a_names(0)
            .iter()
            .any(|s| s == "model.layers.0.linear_attn.in_proj_a.weight"));
        assert!(linear_out_norm_names(0)
            .iter()
            .any(|s| s == "model.layers.0.linear_attn.norm.weight"));
    }
}