fuel-transformers 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! BigCode (StarCoder-1) decoder ported to the lazy-graph API.
//!
//! Phase D LLM port. BigCode is a GPT-2-style transformer:
//! - **Learned absolute position embeddings** (`wpe`) added to the
//!   token embedding (`wte`) at the start. No RoPE, no ALiBi.
//! - **Multi-query attention** — single shared K and V across all
//!   attention heads (`multi_query == true` by default).
//! - LayerNorm with bias on input + post-attention paths.
//! - **GELU MLP** — `down(gelu_tanh(up(x)))`, no gate path.
//! - Q/K/V/O and MLP projections all have biases.
//!
//! The learned-position path is the only thing fundamentally
//! different from StarCoder2 — replace the RoPE step with a single
//! `index_select(wpe, [start_pos..start_pos+seq])` + `broadcast_add`
//! to the token embedding.

use fuel_ir::Shape;
use fuel_ir::error::Result;
use fuel_tensor::Device;
use fuel_tensor::{
    lazy::Tensor, lazy::WeightStorage, lazy::load_tensor_as_f32,
    lazy::load_transposed_matrix_preserve_dtype,
};
use std::sync::Arc;

#[derive(Debug, Clone, PartialEq)]
pub struct BigCodeConfig {
    pub vocab_size: usize,
    pub max_position_embeddings: usize,
    pub num_hidden_layers: usize,
    pub hidden_size: usize,
    pub num_attention_heads: usize,
    pub layer_norm_epsilon: f64,
    pub intermediate_size: usize,
    /// `true` for StarCoder-1 — single shared K/V across all
    /// attention heads.
    pub multi_query: bool,
}

impl BigCodeConfig {
    pub fn head_dim(&self) -> usize {
        self.hidden_size / self.num_attention_heads
    }

    pub fn kv_dim(&self) -> usize {
        if self.multi_query {
            self.head_dim()
        } else {
            self.hidden_size
        }
    }

    /// StarCoder-1 ~1B preset (HuggingFace `bigcode/starcoder`).
    pub fn starcoder_1b() -> Self {
        Self {
            vocab_size: 49_152,
            max_position_embeddings: 8192,
            num_hidden_layers: 24,
            hidden_size: 2048,
            num_attention_heads: 16,
            layer_norm_epsilon: 1e-5,
            intermediate_size: 8192,
            multi_query: true,
        }
    }
}

#[derive(Debug, Clone)]
pub struct BigCodeLayerWeights {
    pub input_ln_gain: Arc<[f32]>,
    pub input_ln_bias: Arc<[f32]>,
    pub post_attn_ln_gain: Arc<[f32]>,
    pub post_attn_ln_bias: Arc<[f32]>,
    pub attn_q: WeightStorage,
    pub attn_q_bias: Arc<[f32]>,
    pub attn_k: WeightStorage,
    pub attn_k_bias: Arc<[f32]>,
    pub attn_v: WeightStorage,
    pub attn_v_bias: Arc<[f32]>,
    pub attn_o: WeightStorage,
    pub attn_o_bias: Arc<[f32]>,
    pub mlp_fc: WeightStorage,
    pub mlp_fc_bias: Arc<[f32]>,
    pub mlp_proj: WeightStorage,
    pub mlp_proj_bias: Arc<[f32]>,
}

#[derive(Debug, Clone)]
pub struct BigCodeWeights {
    /// `[vocab_size, hidden_size]` — wte.
    pub token_embedding: Arc<[f32]>,
    /// `[max_position_embeddings, hidden_size]` — wpe (learned
    /// positional embedding).
    pub position_embedding: Arc<[f32]>,
    pub layers: Vec<BigCodeLayerWeights>,
    pub final_ln_gain: Arc<[f32]>,
    pub final_ln_bias: Arc<[f32]>,
    /// Often tied to `token_embedding` in BigCode checkpoints; the
    /// safetensors loader handles the tying.
    pub output: WeightStorage,
}

#[derive(Debug, Clone)]
pub struct BigCodeModel {
    pub config: BigCodeConfig,
    pub weights: BigCodeWeights,
}

impl BigCodeModel {
    pub fn forward(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let h_norm = self.run_backbone(tokens, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Run the decoder forward up to the final LayerNorm and
    /// return per-token hidden states `(1, seq, hidden_size)`.
    /// Skips the `lm_head` projection. BigCode (StarCoder1)
    /// uses learned absolute position embeddings + LayerNorm
    /// final norm — same backbone is run for both `forward`
    /// and `forward_hidden`.
    pub fn forward_hidden(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        self.run_backbone(tokens, start_pos)
    }

    /// Multimodal entry point. Skips token embedding; runs the decoder
    /// over pre-embedded inputs. Positional encoding (BigCode's learned
    /// wpe table) is still applied internally based on `start_pos` +
    /// seq from the embeds — multimodal hosts pass the token-level
    /// embeds (post token-table lookup, pre-positional-add).
    pub fn forward_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let h_norm = self.run_backbone_embeds(embeds, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Hidden-state variant of [`Self::forward_embeds`].
    pub fn forward_hidden_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        self.run_backbone_embeds(embeds, start_pos)
    }

    /// Build per-token embeddings without running the decoder. NOTE:
    /// the positional embedding is NOT added — callers must pass the
    /// raw embeds to [`Self::forward_embeds`] which adds positional
    /// encoding internally.
    pub fn embed_tokens_anchored(&self, anchor: &Tensor, tokens: &[u32]) -> Result<Tensor> {
        let cfg = &self.config;
        anchor.embed_tokens_anchored(
            self.weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
        )
    }

    fn apply_lm_head(&self, h_norm: &Tensor) -> Result<Tensor> {
        let cfg = &self.config;
        self.weights
            .output
            .apply_linear(h_norm, cfg.hidden_size, cfg.vocab_size)
    }

    fn run_backbone(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let seq = tokens.len();
        assert!(seq > 0);

        let h = Tensor::embed_tokens(
            weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
            &Device::cpu(),
        )?;
        self.run_backbone_embeds(&h, start_pos)
    }

    fn run_backbone_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let dims = embeds.shape();
        let dims = dims.dims();
        if dims.len() != 3 || dims[0] != 1 || dims[2] != cfg.hidden_size {
            return Err(fuel_ir::error::Error::Msg(format!(
                "BigCodeModel::forward_embeds: expected embeds shape (1, seq, hidden_size={}), got {:?}",
                cfg.hidden_size, dims,
            )).bt());
        }
        let seq = dims[1];
        let batch = dims[0];
        if seq == 0 {
            return Err(fuel_ir::error::Error::Msg(
                "BigCodeModel::forward_embeds: seq must be > 0".into(),
            )
            .bt());
        }
        if cfg.num_attention_heads * cfg.head_dim() != cfg.hidden_size {
            return Err(fuel_ir::error::Error::Msg(
                "BigCodeConfig: num_attention_heads * head_dim must equal hidden_size".into(),
            )
            .bt());
        }
        if start_pos + seq > cfg.max_position_embeddings {
            return Err(fuel_ir::error::Error::Msg(format!(
                "BigCodeModel: start_pos + seq ({}) exceeds max_position_embeddings ({})",
                start_pos + seq,
                cfg.max_position_embeddings,
            ))
            .bt());
        }

        let wpe = embeds.const_f32_like(
            weights.position_embedding.clone(),
            Shape::from_dims(&[cfg.max_position_embeddings, cfg.hidden_size]),
        )?;
        let pos_ids: Vec<u32> = (0..seq).map(|i| (start_pos + i) as u32).collect();
        let pos_ids_t = embeds.const_u32_like(pos_ids, Shape::from_dims(&[seq]))?;
        let pos_emb = wpe
            .index_select(0_usize, &pos_ids_t)?
            .reshape(Shape::from_dims(&[batch, seq, cfg.hidden_size]))?;
        let mut h = embeds.add(&pos_emb)?;

        for layer in &weights.layers {
            h = self.apply_layer(&h, layer)?;
        }
        h.layer_norm_affine(
            std::sync::Arc::clone(&weights.final_ln_gain),
            std::sync::Arc::clone(&weights.final_ln_bias),
            cfg.layer_norm_epsilon,
        )
    }

    fn apply_layer(&self, x: &Tensor, layer: &BigCodeLayerWeights) -> Result<Tensor> {
        let cfg = &self.config;
        let head_dim = cfg.head_dim();
        let kv_dim = cfg.kv_dim();
        let x_shape = x.shape();
        let dims = x_shape.dims();
        let batch = dims[0];
        let seq = dims[1];

        let x_norm = x.layer_norm_affine(
            std::sync::Arc::clone(&layer.input_ln_gain),
            std::sync::Arc::clone(&layer.input_ln_bias),
            cfg.layer_norm_epsilon,
        )?;

        let q = layer.attn_q.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            cfg.hidden_size,
            std::sync::Arc::clone(&layer.attn_q_bias),
        )?;
        let k = layer.attn_k.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            kv_dim,
            std::sync::Arc::clone(&layer.attn_k_bias),
        )?;
        let v = layer.attn_v.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            kv_dim,
            std::sync::Arc::clone(&layer.attn_v_bias),
        )?;

        let n_kv_heads = if cfg.multi_query {
            1
        } else {
            cfg.num_attention_heads
        };
        let _ = (batch, seq);
        let q = q.split_heads(cfg.num_attention_heads, head_dim)?;
        let k = k.split_heads(n_kv_heads, head_dim)?;
        let v = v.split_heads(n_kv_heads, head_dim)?;

        // MQA replication: broadcast K and V from 1 → num_heads.
        let n_rep = cfg.num_attention_heads / n_kv_heads;
        let k_full = k.repeat_interleave(1_usize, n_rep)?;
        let v_full = v.repeat_interleave(1_usize, n_rep)?;

        let k_t = k_full.transpose()?;
        let scale = 1.0_f64 / (head_dim as f64).sqrt();
        let scores = q.matmul(&k_t)?;
        let scores_scaled = scores.mul_scalar(scale);
        let mask = Tensor::additive_causal_mask_like(x, seq)
            .reshape(Shape::from_dims(&[1, 1, seq, seq]))?;
        let scores_masked = scores_scaled.broadcast_add(&mask)?;
        let attn = scores_masked.softmax_last_dim()?;
        let attn_v = attn.matmul(&v_full)?;

        let merged = attn_v.merge_heads()?;
        let attn_out = layer.attn_o.apply_linear_with_bias(
            &merged,
            cfg.hidden_size,
            cfg.hidden_size,
            std::sync::Arc::clone(&layer.attn_o_bias),
        )?;

        let h1 = x.add(&attn_out)?;
        let h1_norm = h1.layer_norm_affine(
            std::sync::Arc::clone(&layer.post_attn_ln_gain),
            std::sync::Arc::clone(&layer.post_attn_ln_bias),
            cfg.layer_norm_epsilon,
        )?;

        // GELU MLP.
        let mid = layer.mlp_fc.apply_linear_with_bias(
            &h1_norm,
            cfg.hidden_size,
            cfg.intermediate_size,
            std::sync::Arc::clone(&layer.mlp_fc_bias),
        )?;
        let mid_act = mid.gelu();
        let ffn_out = layer.mlp_proj.apply_linear_with_bias(
            &mid_act,
            cfg.intermediate_size,
            cfg.hidden_size,
            std::sync::Arc::clone(&layer.mlp_proj_bias),
        )?;

        h1.add(&ffn_out)
    }
}

// ---- Safetensors loader ----------------------------------------------------

impl BigCodeWeights {
    /// Load BigCode (GPT-BigCode / StarCoder1) weights from a
    /// `MmapedSafetensors` file using the standard HuggingFace
    /// naming convention.
    ///
    /// BigCode's defining quirk is the fused QKV projection — every
    /// layer's `transformer.h.{i}.attn.c_attn.weight` is a single
    /// `[hidden + 2*kv_dim, hidden]` matrix (Q + K + V concatenated
    /// along the output axis, HF `[out, in]` layout), with a matching
    /// `[hidden + 2*kv_dim]` bias. This loader splits the fused tensor
    /// into separate Q/K/V `WeightStorage` entries (and biases) so the
    /// lazy decoder can reuse the same per-projection apply path it
    /// uses elsewhere.
    ///
    /// Expected names:
    /// - `transformer.wte.weight` → token embedding `[vocab, hidden]`
    /// - `transformer.wpe.weight` → learned position embedding
    ///   `[max_position_embeddings, hidden]`
    /// - `transformer.h.{i}.ln_1.{weight,bias}` → input LayerNorm
    /// - `transformer.h.{i}.attn.c_attn.{weight,bias}` → fused QKV
    ///   (split here)
    /// - `transformer.h.{i}.attn.c_proj.{weight,bias}` → attention output
    /// - `transformer.h.{i}.ln_2.{weight,bias}` → post-attention LN
    /// - `transformer.h.{i}.mlp.c_fc.{weight,bias}` → MLP fc1
    ///   (`[intermediate, hidden]` HF)
    /// - `transformer.h.{i}.mlp.c_proj.{weight,bias}` → MLP fc2
    ///   (`[hidden, intermediate]` HF)
    /// - `transformer.ln_f.{weight,bias}` → final LayerNorm
    /// - `lm_head.weight` (optional) — falls back to tied embedding
    ///   when absent (the GPT-BigCode reference ties `lm_head` to
    ///   `transformer.wte`).
    pub fn load_from_mmapped(
        st: &fuel_loaders::safetensors::MmapedSafetensors,
        cfg: &BigCodeConfig,
    ) -> Result<Self> {
        let h = cfg.hidden_size;
        let kv = cfg.kv_dim();
        let i_dim = cfg.intermediate_size;
        let qkv_out = h + 2 * kv;

        let token_embedding = load_tensor_as_f32(st, "transformer.wte.weight")?;
        if token_embedding.len() != cfg.vocab_size * h {
            fuel_ir::bail!(
                "transformer.wte.weight: {} elts, expected {} ({}×{})",
                token_embedding.len(),
                cfg.vocab_size * h,
                cfg.vocab_size,
                h,
            );
        }

        let position_embedding = load_tensor_as_f32(st, "transformer.wpe.weight")?;
        if position_embedding.len() != cfg.max_position_embeddings * h {
            fuel_ir::bail!(
                "transformer.wpe.weight: {} elts, expected {} ({}×{})",
                position_embedding.len(),
                cfg.max_position_embeddings * h,
                cfg.max_position_embeddings,
                h,
            );
        }

        let mut layers: Vec<BigCodeLayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        for li in 0..cfg.num_hidden_layers {
            let p = format!("transformer.h.{li}");

            // Input + post-attention LayerNorm gain + bias.
            let input_ln_gain = load_tensor_as_f32(st, &format!("{p}.ln_1.weight"))?;
            let input_ln_bias = load_tensor_as_f32(st, &format!("{p}.ln_1.bias"))?;
            let post_attn_ln_gain = load_tensor_as_f32(st, &format!("{p}.ln_2.weight"))?;
            let post_attn_ln_bias = load_tensor_as_f32(st, &format!("{p}.ln_2.bias"))?;

            // Fused QKV weight: HF stores `[hidden + 2*kv_dim, hidden]`
            // (out, in). Read as f32, split into three sub-matrices
            // along the output axis, and physically transpose each to
            // fuel's `[in, out]` layout before storing.
            let fused = load_tensor_as_f32(st, &format!("{p}.attn.c_attn.weight"))?;
            if fused.len() != qkv_out * h {
                fuel_ir::bail!(
                    "{p}.attn.c_attn.weight: {} elts, expected {} ({}×{})",
                    fused.len(),
                    qkv_out * h,
                    qkv_out,
                    h,
                );
            }
            let (q_buf, k_buf, v_buf) = split_fused_qkv_weight(&fused, h, kv, h);
            let attn_q = WeightStorage::F32(Arc::from(q_buf));
            let attn_k = WeightStorage::F32(Arc::from(k_buf));
            let attn_v = WeightStorage::F32(Arc::from(v_buf));

            // Fused QKV bias: `[hidden + 2*kv_dim]`. Split into Q / K / V.
            let fused_bias = load_tensor_as_f32(st, &format!("{p}.attn.c_attn.bias"))?;
            if fused_bias.len() != qkv_out {
                fuel_ir::bail!(
                    "{p}.attn.c_attn.bias: {} elts, expected {}",
                    fused_bias.len(),
                    qkv_out,
                );
            }
            let attn_q_bias: Arc<[f32]> = Arc::from(&fused_bias[..h]);
            let attn_k_bias: Arc<[f32]> = Arc::from(&fused_bias[h..h + kv]);
            let attn_v_bias: Arc<[f32]> = Arc::from(&fused_bias[h + kv..h + 2 * kv]);

            // Attention output projection — standard `[out, in]` HF
            // layout, dtype preserved.
            let attn_o = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.attn.c_proj.weight"),
                h,
                h,
            )?;
            let attn_o_bias_v = load_tensor_as_f32(st, &format!("{p}.attn.c_proj.bias"))?;
            if attn_o_bias_v.len() != h {
                fuel_ir::bail!(
                    "{p}.attn.c_proj.bias: {} elts, expected {}",
                    attn_o_bias_v.len(),
                    h,
                );
            }
            let attn_o_bias: Arc<[f32]> = Arc::from(attn_o_bias_v);

            // MLP fc + projection.
            let mlp_fc = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.c_fc.weight"),
                i_dim,
                h,
            )?;
            let mlp_fc_bias_v = load_tensor_as_f32(st, &format!("{p}.mlp.c_fc.bias"))?;
            if mlp_fc_bias_v.len() != i_dim {
                fuel_ir::bail!(
                    "{p}.mlp.c_fc.bias: {} elts, expected {}",
                    mlp_fc_bias_v.len(),
                    i_dim,
                );
            }
            let mlp_fc_bias: Arc<[f32]> = Arc::from(mlp_fc_bias_v);

            let mlp_proj = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.c_proj.weight"),
                h,
                i_dim,
            )?;
            let mlp_proj_bias_v = load_tensor_as_f32(st, &format!("{p}.mlp.c_proj.bias"))?;
            if mlp_proj_bias_v.len() != h {
                fuel_ir::bail!(
                    "{p}.mlp.c_proj.bias: {} elts, expected {}",
                    mlp_proj_bias_v.len(),
                    h,
                );
            }
            let mlp_proj_bias: Arc<[f32]> = Arc::from(mlp_proj_bias_v);

            layers.push(BigCodeLayerWeights {
                input_ln_gain: Arc::from(input_ln_gain),
                input_ln_bias: Arc::from(input_ln_bias),
                post_attn_ln_gain: Arc::from(post_attn_ln_gain),
                post_attn_ln_bias: Arc::from(post_attn_ln_bias),
                attn_q,
                attn_q_bias,
                attn_k,
                attn_k_bias,
                attn_v,
                attn_v_bias,
                attn_o,
                attn_o_bias,
                mlp_fc,
                mlp_fc_bias,
                mlp_proj,
                mlp_proj_bias,
            });
        }

        let final_ln_gain = load_tensor_as_f32(st, "transformer.ln_f.weight")?;
        let final_ln_bias = load_tensor_as_f32(st, "transformer.ln_f.bias")?;

        // GPT-BigCode reference ties `lm_head` to `transformer.wte`
        // (see `fuel_transformers::models::llm::bigcode`). Try to
        // load an explicit `lm_head.weight` first; if absent, fall
        // back to a transposed copy of the token embedding.
        let output: WeightStorage =
            match load_transposed_matrix_preserve_dtype(st, "lm_head.weight", cfg.vocab_size, h) {
                Ok(w) => w,
                Err(_) => {
                    let mut transposed = vec![0.0_f32; h * cfg.vocab_size];
                    for i in 0..cfg.vocab_size {
                        for j in 0..h {
                            transposed[j * cfg.vocab_size + i] = token_embedding[i * h + j];
                        }
                    }
                    WeightStorage::F32(Arc::from(transposed))
                }
            };

        Ok(BigCodeWeights {
            token_embedding: Arc::from(token_embedding),
            position_embedding: Arc::from(position_embedding),
            layers,
            final_ln_gain: Arc::from(final_ln_gain),
            final_ln_bias: Arc::from(final_ln_bias),
            output,
        })
    }
}

/// Split a fused QKV weight `[q_out + k_out + v_out, in]` (HF row-major
/// `[out, in]`) into three fuel-layout `[in, out_i]` buffers — one each
/// for Q, K, V. Each output buffer is the physical transpose of the
/// corresponding row slice of `fused`.
fn split_fused_qkv_weight(
    fused: &[f32],
    q_out: usize,
    kv_out: usize,
    in_dim: usize,
) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
    let mut q = vec![0.0_f32; in_dim * q_out];
    let mut k = vec![0.0_f32; in_dim * kv_out];
    let mut v = vec![0.0_f32; in_dim * kv_out];
    // Q rows: 0 .. q_out
    for i in 0..q_out {
        for j in 0..in_dim {
            q[j * q_out + i] = fused[i * in_dim + j];
        }
    }
    // K rows: q_out .. q_out + kv_out
    for i in 0..kv_out {
        for j in 0..in_dim {
            k[j * kv_out + i] = fused[(q_out + i) * in_dim + j];
        }
    }
    // V rows: q_out + kv_out .. q_out + 2*kv_out
    for i in 0..kv_out {
        for j in 0..in_dim {
            v[j * kv_out + i] = fused[(q_out + kv_out + i) * in_dim + j];
        }
    }
    (q, k, v)
}

#[cfg(test)]
mod tests {
    use super::*;
    fn tiny_weights(cfg: &BigCodeConfig) -> BigCodeWeights {
        let mut s: u32 = 12121;
        let next = || -> f32 {
            s = s.wrapping_mul(1103515245).wrapping_add(12345);
            ((s >> 16) as u16 as f32 / 65535.0 - 0.5) * 0.05
        };
        let vec_of = |n: usize, next: &mut dyn FnMut() -> f32| -> Arc<[f32]> {
            Arc::from((0..n).map(|_| next()).collect::<Vec<_>>())
        };
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let kv = cfg.kv_dim();
        let mut nb: Box<dyn FnMut() -> f32> = Box::new(next);
        let token_embedding = vec_of(cfg.vocab_size * h, &mut *nb);
        let position_embedding = vec_of(cfg.max_position_embeddings * h, &mut *nb);
        let layers: Vec<BigCodeLayerWeights> = (0..cfg.num_hidden_layers)
            .map(|_| BigCodeLayerWeights {
                input_ln_gain: Arc::from(vec![1.0_f32; h]),
                input_ln_bias: Arc::from(vec![0.0_f32; h]),
                post_attn_ln_gain: Arc::from(vec![1.0_f32; h]),
                post_attn_ln_bias: Arc::from(vec![0.0_f32; h]),
                attn_q: WeightStorage::F32(vec_of(h * h, &mut *nb)),
                attn_q_bias: vec_of(h, &mut *nb),
                attn_k: WeightStorage::F32(vec_of(h * kv, &mut *nb)),
                attn_k_bias: vec_of(kv, &mut *nb),
                attn_v: WeightStorage::F32(vec_of(h * kv, &mut *nb)),
                attn_v_bias: vec_of(kv, &mut *nb),
                attn_o: WeightStorage::F32(vec_of(h * h, &mut *nb)),
                attn_o_bias: vec_of(h, &mut *nb),
                mlp_fc: WeightStorage::F32(vec_of(h * i, &mut *nb)),
                mlp_fc_bias: vec_of(i, &mut *nb),
                mlp_proj: WeightStorage::F32(vec_of(i * h, &mut *nb)),
                mlp_proj_bias: vec_of(h, &mut *nb),
            })
            .collect();
        let final_ln_gain = Arc::from(vec![1.0_f32; h]);
        let final_ln_bias = Arc::from(vec![0.0_f32; h]);
        let output = WeightStorage::F32(vec_of(h * cfg.vocab_size, &mut *nb));
        BigCodeWeights {
            token_embedding,
            position_embedding,
            layers,
            final_ln_gain,
            final_ln_bias,
            output,
        }
    }

    #[test]
    fn forward_shape_and_finite_mqa() {
        let cfg = BigCodeConfig {
            vocab_size: 32,
            max_position_embeddings: 64,
            num_hidden_layers: 2,
            hidden_size: 16,
            num_attention_heads: 4,
            layer_norm_epsilon: 1e-5,
            intermediate_size: 32,
            multi_query: true,
        };
        let model = BigCodeModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let logits = model.forward(&[1, 2, 3], 0).unwrap();
        assert_eq!(logits.shape().dims(), &[1, 3, cfg.vocab_size]);
        for &v in &logits.realize_f32() {
            assert!(v.is_finite());
        }
    }

    /// Different start_pos must produce different output (the
    /// learned position embedding pulls different rows).
    #[test]
    fn different_start_pos_changes_output() {
        let cfg = BigCodeConfig {
            vocab_size: 16,
            max_position_embeddings: 32,
            num_hidden_layers: 1,
            hidden_size: 8,
            num_attention_heads: 2,
            layer_norm_epsilon: 1e-5,
            intermediate_size: 16,
            multi_query: true,
        };
        let weights = tiny_weights(&cfg);
        let out_0 = BigCodeModel {
            config: cfg.clone(),
            weights: weights.clone(),
        }
        .forward(&[1, 2, 3], 0)
        .unwrap()
        .realize_f32();
        let out_5 = BigCodeModel {
            config: cfg,
            weights,
        }
        .forward(&[1, 2, 3], 5)
        .unwrap()
        .realize_f32();
        let any_diff = out_0
            .iter()
            .zip(out_5.iter())
            .any(|(&a, &b)| (a - b).abs() > 1e-7);
        assert!(
            any_diff,
            "different start_pos must change the learned-position output"
        );
    }

    /// `forward_hidden` returns post-LayerNorm hidden states.
    #[test]
    fn forward_hidden_shape_and_finite() {
        let cfg = BigCodeConfig {
            vocab_size: 32,
            max_position_embeddings: 64,
            num_hidden_layers: 2,
            hidden_size: 16,
            num_attention_heads: 4,
            layer_norm_epsilon: 1e-5,
            intermediate_size: 32,
            multi_query: true,
        };
        let model = BigCodeModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3, 4];
        let hidden = model.forward_hidden(&tokens, 0).unwrap();
        assert_eq!(hidden.shape().dims(), &[1, tokens.len(), cfg.hidden_size]);
        for &v in &hidden.realize_f32() {
            assert!(v.is_finite(), "non-finite hidden: {v}");
        }
    }

    fn forward_embeds_test_cfg() -> BigCodeConfig {
        BigCodeConfig {
            vocab_size: 32,
            max_position_embeddings: 64,
            num_hidden_layers: 2,
            hidden_size: 16,
            num_attention_heads: 4,
            layer_norm_epsilon: 1e-5,
            intermediate_size: 32,
            multi_query: true,
        }
    }

    #[test]
    fn forward_embeds_matches_forward_after_token_lookup() {
        let cfg = forward_embeds_test_cfg();
        let model = BigCodeModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3];
        let logits_ref = model.forward(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let logits_via_embeds = model.forward_embeds(&embeds, 0).unwrap().realize_f32();
        let max_diff = logits_ref
            .iter()
            .zip(logits_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-5,
            "BigCode forward vs forward_embeds must agree (max diff {max_diff})"
        );
    }

    #[test]
    fn forward_embeds_rejects_bad_shape() {
        let cfg = forward_embeds_test_cfg();
        let model = BigCodeModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let bad = Tensor::from_f32(
            vec![0.0_f32; 3 * (cfg.hidden_size + 1)],
            Shape::from_dims(&[1, 3, cfg.hidden_size + 1]),
            &Device::cpu(),
        )
        .unwrap();
        assert!(model.forward_embeds(&bad, 0).is_err());
    }

    #[test]
    fn forward_hidden_embeds_matches_forward_hidden() {
        let cfg = forward_embeds_test_cfg();
        let model = BigCodeModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![5, 7];
        let h_ref = model.forward_hidden(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let h_via_embeds = model
            .forward_hidden_embeds(&embeds, 0)
            .unwrap()
            .realize_f32();
        let max_diff = h_ref
            .iter()
            .zip(h_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-5,
            "BigCode forward_hidden vs forward_hidden_embeds must agree (max diff {max_diff})"
        );
    }
}