From b94ed59f38a8fc5839e684e56bd78bc9e8dd2b7c Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:27 -0400 Subject: [PATCH 01/24] fix: check matrix and dequantization size arithmetic --- src/ops.rs | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/ops.rs b/src/ops.rs index 2b9f754..4e82c5e 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -10,11 +10,15 @@ //! * `gelu` uses the tanh approximation, matching `jax.nn.gelu(approximate=True)`. //! * `rope` uses the rotate-half (GPT-NeoX / HF) convention. +fn checked_elements(rows: usize, cols: usize) -> usize { + rows.checked_mul(cols).expect("operator dimensions overflow") +} + /// Standard matrix multiply: `A`[m×k] · `B`[k×n] → `C`[m×n] (row-major). pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { - assert_eq!(a.len(), m * k, "A has wrong length"); - assert_eq!(b.len(), k * n, "B has wrong length"); - let mut c = vec![0.0f32; m * n]; + assert_eq!(a.len(), checked_elements(m, k), "A has wrong length"); + assert_eq!(b.len(), checked_elements(k, n), "B has wrong length"); + let mut c = vec![0.0f32; checked_elements(m, n)]; for row in 0..m { for i in 0..k { let a_ik = a[row * k + i]; @@ -34,10 +38,10 @@ pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { /// /// `x`[n×in], `w`[out×in], `b`[out] → `[n×out]`. pub fn linear(x: &[f32], w: &[f32], b: &[f32], n: usize, in_f: usize, out_f: usize) -> Vec { - assert_eq!(x.len(), n * in_f); - assert_eq!(w.len(), out_f * in_f); + assert_eq!(x.len(), checked_elements(n, in_f)); + assert_eq!(w.len(), checked_elements(out_f, in_f)); assert_eq!(b.len(), out_f); - let mut y = vec![0.0f32; n * out_f]; + let mut y = vec![0.0f32; checked_elements(n, out_f)]; for row in 0..n { for o in 0..out_f { let mut acc = b[o]; @@ -234,9 +238,9 @@ pub fn mha(q: &[f32], k: &[f32], v: &[f32], seq: usize, heads: usize, head_dim: /// Dequantizes a per-row INT8 weight matrix: `out[r,c] = q[r,c] · scale[r]`. pub fn dequantize_int8(q: &[i8], scales: &[f32], rows: usize, cols: usize) -> Vec { - assert_eq!(q.len(), rows * cols); + assert_eq!(q.len(), checked_elements(rows, cols)); assert_eq!(scales.len(), rows); - let mut out = vec![0.0f32; rows * cols]; + let mut out = vec![0.0f32; checked_elements(rows, cols)]; for r in 0..rows { let s = scales[r]; for c in 0..cols { @@ -369,4 +373,10 @@ mod tests { let out = dequantize_int8(&q, &scales, 1, 3); approx(&out, &[1.0, -1.0, 10.0], 1e-6); } + + #[test] + #[should_panic(expected = "operator dimensions overflow")] + fn matmul_rejects_overflow_before_allocation() { + matmul(&[], &[], usize::MAX, 2, 1); + } } From a5034824993df83b9049dffef706dbaf441e0a47 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:28 -0400 Subject: [PATCH 02/24] fix: reject invalid normalization epsilon --- src/ops.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/ops.rs b/src/ops.rs index 4e82c5e..2fb49e9 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -99,6 +99,7 @@ pub fn layernorm( d: usize, eps: f32, ) -> Vec { + assert!(eps.is_finite() && eps >= 0.0, "normalization epsilon must be finite and non-negative"); assert_eq!(x.len(), rows * d); assert_eq!(gamma.len(), d); assert_eq!(beta.len(), d); @@ -117,6 +118,7 @@ pub fn layernorm( /// Row-wise RMSNorm over the last dimension of size `d`: `y = x / sqrt(mean(x²) + eps) * gamma`. pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec { + assert!(eps.is_finite() && eps >= 0.0, "normalization epsilon must be finite and non-negative"); assert_eq!(x.len(), rows * d); assert_eq!(gamma.len(), d); let mut out = vec![0.0f32; rows * d]; @@ -379,4 +381,12 @@ mod tests { fn matmul_rejects_overflow_before_allocation() { matmul(&[], &[], usize::MAX, 2, 1); } + + #[test] + fn normalization_rejects_invalid_epsilon() { + for eps in [f32::NAN, f32::INFINITY, -1.0] { + assert!(std::panic::catch_unwind(|| layernorm(&[1.0], &[1.0], &[0.0], 1, 1, eps)).is_err()); + assert!(std::panic::catch_unwind(|| rmsnorm(&[1.0], &[1.0], 1, 1, eps)).is_err()); + } + } } From 0b8f754f13ed796adeb1c024bc2bbd58e8d35398 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:29 -0400 Subject: [PATCH 03/24] fix: validate normalization width and buffer sizes --- src/ops.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/ops.rs b/src/ops.rs index 2fb49e9..ba422f0 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -99,11 +99,12 @@ pub fn layernorm( d: usize, eps: f32, ) -> Vec { + assert!(d > 0, "normalization width must be positive"); assert!(eps.is_finite() && eps >= 0.0, "normalization epsilon must be finite and non-negative"); - assert_eq!(x.len(), rows * d); + assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(gamma.len(), d); assert_eq!(beta.len(), d); - let mut out = vec![0.0f32; rows * d]; + let mut out = vec![0.0f32; checked_elements(rows, d)]; for r in 0..rows { let row = &x[r * d..r * d + d]; let mean = row.iter().sum::() / d as f32; @@ -118,10 +119,11 @@ pub fn layernorm( /// Row-wise RMSNorm over the last dimension of size `d`: `y = x / sqrt(mean(x²) + eps) * gamma`. pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec { + assert!(d > 0, "normalization width must be positive"); assert!(eps.is_finite() && eps >= 0.0, "normalization epsilon must be finite and non-negative"); - assert_eq!(x.len(), rows * d); + assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(gamma.len(), d); - let mut out = vec![0.0f32; rows * d]; + let mut out = vec![0.0f32; checked_elements(rows, d)]; for r in 0..rows { let row = &x[r * d..r * d + d]; let ms = row.iter().map(|v| v * v).sum::() / d as f32; @@ -389,4 +391,11 @@ mod tests { assert!(std::panic::catch_unwind(|| rmsnorm(&[1.0], &[1.0], 1, 1, eps)).is_err()); } } + + #[test] + fn normalization_distinguishes_empty_batches_from_zero_width() { + assert!(layernorm(&[], &[1.0], &[0.0], 0, 1, 1e-5).is_empty()); + assert!(rmsnorm(&[], &[1.0], 0, 1, 0.0).is_empty()); + assert!(std::panic::catch_unwind(|| layernorm(&[], &[], &[], 1, 0, 0.0)).is_err()); + } } From 66fb160d2003eebe8ffd8cbeaf86bb2420c0da68 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:31 -0400 Subject: [PATCH 04/24] fix: validate rotary embedding dimensions and frequency base --- src/ops.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/ops.rs b/src/ops.rs index ba422f0..d61177b 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -139,7 +139,9 @@ pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec /// `x`[rows×d] in place. `positions[r]` gives the absolute position of row `r`. /// `d` must be even. pub fn rope_inplace(x: &mut [f32], positions: &[usize], rows: usize, d: usize, theta: f32) { - assert_eq!(x.len(), rows * d); + assert!(d > 0, "rope head dim must be positive"); + assert!(theta.is_finite() && theta > 0.0, "rope theta must be finite and positive"); + assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(positions.len(), rows); assert_eq!(d % 2, 0, "rope head dim must be even"); let half = d / 2; @@ -398,4 +400,11 @@ mod tests { assert!(rmsnorm(&[], &[1.0], 0, 1, 0.0).is_empty()); assert!(std::panic::catch_unwind(|| layernorm(&[], &[], &[], 1, 0, 0.0)).is_err()); } + + #[test] + fn rope_rejects_invalid_frequency_bases() { + for theta in [0.0, -1.0, f32::NAN, f32::INFINITY] { + assert!(std::panic::catch_unwind(|| rope_inplace(&mut [1.0, 2.0], &[0], 1, 2, theta)).is_err()); + } + } } From dbcdc8bff4b811b83b9c42b05c6d415ceb005abd Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:32 -0400 Subject: [PATCH 05/24] fix: bound attention offsets and reject undefined attention inputs --- src/ops.rs | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/ops.rs b/src/ops.rs index d61177b..d2b5e9f 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -176,15 +176,17 @@ pub fn attention( causal: bool, q_offset: usize, ) -> Vec { - assert_eq!(q.len(), m * d); - assert_eq!(k.len(), seq * d); - assert_eq!(v.len(), seq * d); + assert!(d > 0, "attention width must be positive"); + assert!(seq > 0 || m == 0, "non-empty attention queries require keys"); + assert_eq!(q.len(), checked_elements(m, d)); + assert_eq!(k.len(), checked_elements(seq, d)); + assert_eq!(v.len(), checked_elements(seq, d)); let scale = 1.0 / (d as f32).sqrt(); - let mut out = vec![0.0f32; m * d]; + let mut out = vec![0.0f32; checked_elements(m, d)]; let mut scores = vec![0.0f32; seq]; for qi in 0..m { let limit = if causal { - (qi + q_offset + 1).min(seq) + qi.saturating_add(q_offset).saturating_add(1).min(seq) } else { seq }; @@ -407,4 +409,10 @@ mod tests { assert!(std::panic::catch_unwind(|| rope_inplace(&mut [1.0, 2.0], &[0], 1, 2, theta)).is_err()); } } + + #[test] + fn attention_large_offset_attends_to_all_available_keys() { + let out = attention(&[0.0], &[0.0, 0.0], &[2.0, 4.0], 1, 2, 1, true, usize::MAX); + assert_eq!(out, vec![3.0]); + } } From 8d76f874dcdada85e27bfaaf2049c193bc1d91e3 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:34 -0400 Subject: [PATCH 06/24] fix: validate multi-head attention dimension products --- src/ops.rs | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/ops.rs b/src/ops.rs index d2b5e9f..fbd690d 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -213,12 +213,13 @@ pub fn attention( /// heads laid out contiguously per row; output has the same shape. Query `i` /// attends causally over keys `j <= i`. This is the GPT-2 attention reference. pub fn mha(q: &[f32], k: &[f32], v: &[f32], seq: usize, heads: usize, head_dim: usize) -> Vec { - let hd = heads * head_dim; - assert_eq!(q.len(), seq * hd); - assert_eq!(k.len(), seq * hd); - assert_eq!(v.len(), seq * hd); + assert!(heads > 0 && head_dim > 0, "attention heads and head width must be positive"); + let hd = checked_elements(heads, head_dim); + assert_eq!(q.len(), checked_elements(seq, hd)); + assert_eq!(k.len(), checked_elements(seq, hd)); + assert_eq!(v.len(), checked_elements(seq, hd)); let scale = 1.0 / (head_dim as f32).sqrt(); - let mut out = vec![0.0f32; seq * hd]; + let mut out = vec![0.0f32; checked_elements(seq, hd)]; let mut scores = vec![0.0f32; seq]; for h in 0..heads { let base = h * head_dim; @@ -415,4 +416,11 @@ mod tests { let out = attention(&[0.0], &[0.0, 0.0], &[2.0, 4.0], 1, 2, 1, true, usize::MAX); assert_eq!(out, vec![3.0]); } + + #[test] + fn mha_rejects_zero_and_overflowing_head_dimensions() { + assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, 0, 1)).is_err()); + assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, usize::MAX, 2)).is_err()); + assert!(mha(&[], &[], &[], 0, 1, 1).is_empty()); + } } From 993e1dea59ce5891adb39f42a0e7aafca099114c Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:35 -0400 Subject: [PATCH 07/24] fix: preserve GPT-2 whitespace pre-tokenization boundaries --- src/tokenizer.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 65350b1..7afb49b 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -249,6 +249,9 @@ fn pre_tokenize(text: &str) -> Vec { while e < n && chars[e].is_whitespace() { e += 1; } + if e < n && e - i > 1 { + e -= 1; + } out.push(chars[i..e].iter().collect()); i = e; continue; @@ -365,4 +368,12 @@ mod tests { assert_eq!(tokenizer.encode("ab"), vec![256]); assert_eq!(tokenizer.decode(&[256]), "ab"); } + + #[test] + fn whitespace_runs_match_gpt2_lookahead() { + assert_eq!(pre_tokenize("a b"), vec!["a", " ", " b"]); + assert_eq!(pre_tokenize("a b"), vec!["a", " ", " b"]); + assert_eq!(pre_tokenize("a "), vec!["a", " "]); + assert_eq!(pre_tokenize("a\n b"), vec!["a", "\n", " b"]); + } } From 3468c5c7708726c11d563c48bd6b6a8e94232776 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:36 -0400 Subject: [PATCH 08/24] fix: parse hash-prefixed BPE merges and reject malformed lines --- src/tokenizer.rs | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 7afb49b..8e08993 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -13,6 +13,8 @@ use thiserror::Error; #[derive(Error, Debug)] pub enum TokenizerError { + #[error("invalid BPE merge line: {0:?}")] + InvalidMergeLine(String), #[error("IO error: {0}")] Io(#[from] std::io::Error), #[error("failed to parse vocab.json: {0}")] @@ -96,16 +98,19 @@ impl Tokenizer { let mut bpe_ranks = HashMap::new(); for (rank, line) in merges_raw .lines() - .filter(|l| !l.starts_with('#')) + .map(str::trim) + .filter(|l| !l.is_empty() && !l.starts_with("#version:")) .enumerate() { let mut it = line.split_whitespace(); - if let (Some(a), Some(b)) = (it.next(), it.next()) { + if let (Some(a), Some(b), None) = (it.next(), it.next(), it.next()) { let merged = format!("{a}{b}"); if !vocab.contains_key(&merged) { return Err(TokenizerError::MissingVocabSymbol(merged)); } bpe_ranks.insert((a.to_string(), b.to_string()), rank); + } else { + return Err(TokenizerError::InvalidMergeLine(line.to_string())); } } @@ -376,4 +381,19 @@ mod tests { assert_eq!(pre_tokenize("a "), vec!["a", " "]); assert_eq!(pre_tokenize("a\n b"), vec!["a", "\n", " b"]); } + + #[test] + fn hash_prefix_merges_are_not_comments() { + let mut vocab = byte_vocab(); + vocab.insert("##".into(), 256); + let tokenizer = Tokenizer::from_assets(vocab, "#version: 0.2\n\n# #\n").unwrap(); + assert_eq!(tokenizer.encode("##"), vec![256]); + } + + #[test] + fn malformed_merge_lines_are_reported() { + for line in ["a", "a b c"] { + assert!(matches!(Tokenizer::from_assets(byte_vocab(), line), Err(TokenizerError::InvalidMergeLine(_)))); + } + } } From c6621e04627e75e6b61c4372e5a25997b21bf1c9 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:37 -0400 Subject: [PATCH 09/24] fix: reject duplicate BPE merge pairs --- src/tokenizer.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 8e08993..cde066f 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -13,6 +13,8 @@ use thiserror::Error; #[derive(Error, Debug)] pub enum TokenizerError { + #[error("duplicate BPE merge pair: {0:?}")] + DuplicateMerge(String), #[error("invalid BPE merge line: {0:?}")] InvalidMergeLine(String), #[error("IO error: {0}")] @@ -108,7 +110,9 @@ impl Tokenizer { if !vocab.contains_key(&merged) { return Err(TokenizerError::MissingVocabSymbol(merged)); } - bpe_ranks.insert((a.to_string(), b.to_string()), rank); + if bpe_ranks.insert((a.to_string(), b.to_string()), rank).is_some() { + return Err(TokenizerError::DuplicateMerge(line.to_string())); + } } else { return Err(TokenizerError::InvalidMergeLine(line.to_string())); } @@ -396,4 +400,11 @@ mod tests { assert!(matches!(Tokenizer::from_assets(byte_vocab(), line), Err(TokenizerError::InvalidMergeLine(_)))); } } + + #[test] + fn duplicate_merges_are_rejected() { + let mut vocab = byte_vocab(); + vocab.insert("ab".into(), 256); + assert!(matches!(Tokenizer::from_assets(vocab, "a b\na b\n"), Err(TokenizerError::DuplicateMerge(_)))); + } } From cd7609b848311b11d555ee36950c77baa9ce7bec Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:38 -0400 Subject: [PATCH 10/24] fix: validate BPE merge operand vocabulary entries --- src/tokenizer.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index cde066f..f42705c 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -106,6 +106,11 @@ impl Tokenizer { { let mut it = line.split_whitespace(); if let (Some(a), Some(b), None) = (it.next(), it.next(), it.next()) { + for operand in [a, b] { + if !vocab.contains_key(operand) { + return Err(TokenizerError::MissingVocabSymbol(operand.to_string())); + } + } let merged = format!("{a}{b}"); if !vocab.contains_key(&merged) { return Err(TokenizerError::MissingVocabSymbol(merged)); @@ -407,4 +412,11 @@ mod tests { vocab.insert("ab".into(), 256); assert!(matches!(Tokenizer::from_assets(vocab, "a b\na b\n"), Err(TokenizerError::DuplicateMerge(_)))); } + + #[test] + fn missing_merge_operands_are_reported() { + let mut vocab = byte_vocab(); + vocab.insert("not_presenta".into(), 256); + assert!(matches!(Tokenizer::from_assets(vocab, "not_present a"), Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "not_present")); + } } From e3d280831781c859ae38bce87cde30bb074d0db7 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:40 -0400 Subject: [PATCH 11/24] feat: add checked token decoding APIs --- src/tokenizer.rs | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/tokenizer.rs b/src/tokenizer.rs index f42705c..a051fa2 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -13,6 +13,8 @@ use thiserror::Error; #[derive(Error, Debug)] pub enum TokenizerError { + #[error("unknown tokenizer ID: {0}")] + UnknownTokenId(usize), #[error("duplicate BPE merge pair: {0:?}")] DuplicateMerge(String), #[error("invalid BPE merge line: {0:?}")] @@ -160,6 +162,21 @@ impl Tokenizer { ids } + /// Decodes token IDs, returning an error instead of skipping unknown IDs. + pub fn try_decode(&self, ids: &[usize]) -> Result { + Ok(String::from_utf8_lossy(&self.try_decode_bytes(ids)?).into_owned()) + } + + /// Decodes raw bytes while checking every token ID. + pub fn try_decode_bytes(&self, ids: &[usize]) -> Result, TokenizerError> { + let mut bytes = Vec::new(); + for id in ids { + let token = self.decoder.get(id).ok_or(TokenizerError::UnknownTokenId(*id))?; + bytes.extend(token.chars().filter_map(|c| self.byte_decoder.get(&c).copied())); + } + Ok(bytes) + } + /// Decodes token ids back into text. pub fn decode(&self, ids: &[usize]) -> String { String::from_utf8_lossy(&self.decode_bytes(ids)).into_owned() @@ -419,4 +436,12 @@ mod tests { vocab.insert("not_presenta".into(), 256); assert!(matches!(Tokenizer::from_assets(vocab, "not_present a"), Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "not_present")); } + + #[test] + fn checked_decode_reports_unknown_ids_and_preserves_bytes() { + let tokenizer = Tokenizer::from_assets(byte_vocab(), "").unwrap(); + assert!(matches!(tokenizer.try_decode(&[999]), Err(TokenizerError::UnknownTokenId(999)))); + assert_eq!(tokenizer.try_decode_bytes(&[0xe2, 0x82, 0xac]).unwrap(), vec![0xe2, 0x82, 0xac]); + assert_eq!(tokenizer.try_decode(&[0xe2, 0x82, 0xac]).unwrap(), "€"); + } } From a9eba99ccd880a3466afce7bbb8b4f0ac5c4cf57 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:41 -0400 Subject: [PATCH 12/24] feat: expose checked token sampling and parameter validation --- src/gpt2.rs | 60 +++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 42 insertions(+), 18 deletions(-) diff --git a/src/gpt2.rs b/src/gpt2.rs index 463b722..414c92d 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -19,6 +19,8 @@ use crate::tensor::{Tensor, TensorError}; #[derive(Error, Debug)] pub enum Gpt2Error { + #[error("invalid sampling input: {0}")] + InvalidSampling(&'static str), #[error("tensor '{name}' has shape {found:?}; expected {expected:?}")] BadShape { name: String, @@ -420,18 +422,36 @@ pub struct Sampler { impl Sampler { /// Greedy (argmax) sampling. pub fn greedy() -> Self { - Sampler { - temperature: 0.0, - top_k: 0, - seed: 0, + Sampler { temperature: 0.0, top_k: 0, seed: 0 } + } + + pub fn validate(&self) -> Result<(), Gpt2Error> { + if !self.temperature.is_finite() || self.temperature < 0.0 { + return Err(Gpt2Error::InvalidSampling("temperature must be finite and non-negative")); } + Ok(()) + } + + /// Samples one token using this seed. Generation retains the RNG across tokens. + pub fn sample_token(&self, logits: &[f32]) -> Result { + self.sample_checked(logits, &mut Rng::new(self.seed)) } fn sample(&self, logits: &[f32], rng: &mut Rng) -> usize { - if self.temperature <= 0.0 { - return argmax(logits); + self.sample_checked(logits, rng).expect("invalid token sampling input") + } + + fn sample_checked(&self, logits: &[f32], rng: &mut Rng) -> Result { + self.validate()?; + if logits.is_empty() { + return Err(Gpt2Error::InvalidSampling("logits must not be empty")); + } + if logits.iter().any(|logit| !logit.is_finite()) { + return Err(Gpt2Error::InvalidSampling("logits must be finite")); + } + if self.temperature == 0.0 { + return Ok(argmax(logits)); } - // Optional top-k: keep only the k highest logits. let mut idx: Vec = (0..logits.len()).collect(); if self.top_k > 0 && self.top_k < logits.len() { idx.select_nth_unstable_by(self.top_k, |&a, &b| { @@ -440,23 +460,16 @@ impl Sampler { idx.truncate(self.top_k); } let max = idx.iter().map(|&i| logits[i]).fold(f32::MIN, f32::max); - let mut probs: Vec = idx - .iter() - .map(|&i| ((logits[i] - max) / self.temperature).exp()) - .collect(); + let mut probs: Vec = idx.iter().map(|&i| ((logits[i] - max) / self.temperature).exp()).collect(); let sum: f32 = probs.iter().sum(); - for p in &mut probs { - *p /= sum; - } + for p in &mut probs { *p /= sum; } let r = rng.next_f32(); let mut acc = 0.0; for (j, &p) in probs.iter().enumerate() { acc += p; - if r <= acc { - return idx[j]; - } + if r <= acc { return Ok(idx[j]); } } - idx[idx.len() - 1] + Ok(idx[idx.len() - 1]) } } @@ -588,4 +601,15 @@ mod tests { assert_eq!(calls, 1); assert_eq!(tokens.len(), 2); } + + #[test] + fn checked_sampling_rejects_invalid_inputs() { + for temperature in [f32::NAN, f32::INFINITY, -1.0] { + assert!(Sampler { temperature, top_k: 1, seed: 0 }.sample_token(&[0.0]).is_err()); + } + assert!(Sampler::greedy().sample_token(&[]).is_err()); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert!(Sampler::greedy().sample_token(&[value, 0.0]).is_err()); + } + } } From a3b91a97680625c9f0c19a58224bd2f37c4ee7be Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:42 -0400 Subject: [PATCH 13/24] fix: normalize sampling weights with wider arithmetic --- src/gpt2.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/gpt2.rs b/src/gpt2.rs index 414c92d..c008f9a 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -459,11 +459,11 @@ impl Sampler { }); idx.truncate(self.top_k); } - let max = idx.iter().map(|&i| logits[i]).fold(f32::MIN, f32::max); - let mut probs: Vec = idx.iter().map(|&i| ((logits[i] - max) / self.temperature).exp()).collect(); - let sum: f32 = probs.iter().sum(); + let max = idx.iter().map(|&i| logits[i]).fold(f32::NEG_INFINITY, f32::max); + let mut probs: Vec = idx.iter().map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()).collect(); + let sum: f64 = probs.iter().sum(); for p in &mut probs { *p /= sum; } - let r = rng.next_f32(); + let r = f64::from(rng.next_f32()); let mut acc = 0.0; for (j, &p) in probs.iter().enumerate() { acc += p; @@ -612,4 +612,10 @@ mod tests { assert!(Sampler::greedy().sample_token(&[value, 0.0]).is_err()); } } + + #[test] + fn extreme_finite_logits_sample_without_overflow() { + let sampler = Sampler { temperature: f32::MIN_POSITIVE, top_k: 0, seed: 42 }; + assert_eq!(sampler.sample_token(&[-f32::MAX, f32::MAX]).unwrap(), 1); + } } From c121f85e313b254929916517b302b3a0b0735610 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:44 -0400 Subject: [PATCH 14/24] fix: make top-k tie handling and seeded candidate order deterministic --- src/gpt2.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/gpt2.rs b/src/gpt2.rs index c008f9a..e34cd33 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -455,9 +455,10 @@ impl Sampler { let mut idx: Vec = (0..logits.len()).collect(); if self.top_k > 0 && self.top_k < logits.len() { idx.select_nth_unstable_by(self.top_k, |&a, &b| { - logits[b].partial_cmp(&logits[a]).unwrap() + logits[b].total_cmp(&logits[a]).then_with(|| a.cmp(&b)) }); idx.truncate(self.top_k); + idx.sort_unstable(); } let max = idx.iter().map(|&i| logits[i]).fold(f32::NEG_INFINITY, f32::max); let mut probs: Vec = idx.iter().map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()).collect(); @@ -618,4 +619,14 @@ mod tests { let sampler = Sampler { temperature: f32::MIN_POSITIVE, top_k: 0, seed: 42 }; assert_eq!(sampler.sample_token(&[-f32::MAX, f32::MAX]).unwrap(), 1); } + + #[test] + fn top_k_ties_choose_lowest_ids_and_repeat_from_a_seed() { + let sampler = Sampler { temperature: 1.0, top_k: 1, seed: 42 }; + assert_eq!(sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(), 0); + let sampler = Sampler { top_k: 2, ..sampler }; + let first = sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(); + assert!(first < 2); + assert_eq!(first, sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap()); + } } From f2e7cf6d689c5c351e703ff28fa7a3d66cc9316f Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:45 -0400 Subject: [PATCH 15/24] fix: exclude zero-probability candidates at sampling boundaries --- src/gpt2.rs | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/src/gpt2.rs b/src/gpt2.rs index e34cd33..ac02eb4 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -464,14 +464,19 @@ impl Sampler { let mut probs: Vec = idx.iter().map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()).collect(); let sum: f64 = probs.iter().sum(); for p in &mut probs { *p /= sum; } - let r = f64::from(rng.next_f32()); - let mut acc = 0.0; - for (j, &p) in probs.iter().enumerate() { - acc += p; - if r <= acc { return Ok(idx[j]); } - } - Ok(idx[idx.len() - 1]) + Ok(select_probability(&idx, &probs, f64::from(rng.next_f32()))) + } +} + +fn select_probability(indices: &[usize], probabilities: &[f64], draw: f64) -> usize { + let mut cumulative = 0.0; + for (index, &probability) in probabilities.iter().enumerate() { + cumulative += probability; + if draw < cumulative { return indices[index]; } } + // Rounding near one must still select a candidate with positive weight. + let last = probabilities.iter().rposition(|&p| p > 0.0).expect("sampling has no positive weight"); + indices[last] } fn argmax(v: &[f32]) -> usize { @@ -629,4 +634,10 @@ mod tests { assert!(first < 2); assert_eq!(first, sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap()); } + + #[test] + fn sampling_boundaries_never_choose_zero_weight() { + assert_eq!(select_probability(&[10, 20], &[0.0, 1.0], 0.0), 20); + assert_eq!(select_probability(&[10, 20, 30], &[0.5, 0.4999999999999999, 0.0], 0.9999999999999999), 20); + } } From 29392f127270b1c29e7bfd16c40f0c6c0713724a Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:46 -0400 Subject: [PATCH 16/24] fix: validate generation prompts and stop-token IDs --- src/gpt2.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/gpt2.rs b/src/gpt2.rs index ac02eb4..e74e715 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -368,6 +368,12 @@ impl Gpt2 { mut on_token: impl FnMut(usize) -> bool, ) -> Vec { let mut toks = prompt.to_vec(); + if max_new == 0 { return toks; } + self.config.validate().expect("invalid GPT-2 configuration"); + sampler.validate().expect("invalid generation sampler"); + assert!(!prompt.is_empty(), "generation requires a non-empty prompt"); + assert!(eot_token < self.config.vocab_size, "end-of-text ID exceeds GPT-2 vocabulary"); + assert!(prompt.iter().all(|&token| token < self.config.vocab_size), "prompt ID exceeds GPT-2 vocabulary"); let n_ctx = self.config.n_ctx; let mut rng = Rng::new(sampler.seed); for _ in 0..max_new { @@ -640,4 +646,16 @@ mod tests { assert_eq!(select_probability(&[10, 20], &[0.0, 1.0], 0.0), 20); assert_eq!(select_probability(&[10, 20, 30], &[0.5, 0.4999999999999999, 0.0], 0.9999999999999999), 20); } + + #[test] + fn zero_token_generation_returns_the_prompt_without_inference() { + let tokens = forward_test_model().generate(&crate::model::CpuBackend, &[], 0, &Sampler::greedy(), 1, |_| panic!("unexpected inference")); + assert!(tokens.is_empty()); + } + + #[test] + #[should_panic(expected = "end-of-text ID exceeds")] + fn generation_rejects_unknown_stop_tokens() { + forward_test_model().generate(&crate::model::CpuBackend, &[0], 1, &Sampler::greedy(), 999, |_| {}); + } } From 0ed6e268b8c72bb299cf8a0bdc790e3468662daf Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:47 -0400 Subject: [PATCH 17/24] fix: reject zero-width MLP checkpoint layers --- src/model.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/model.rs b/src/model.rs index 1046f7a..52750ee 100644 --- a/src/model.rs +++ b/src/model.rs @@ -102,7 +102,7 @@ impl Mlp { .ok_or_else(|| ModelError::MissingTensor(b_name.clone()))?; w.validate()?; b.validate()?; - if w.shape.len() != 2 { + if w.shape.len() != 2 || w.shape.contains(&0) { return Err(ModelError::BadShape { name: w_name, shape: w.shape.clone(), @@ -264,4 +264,12 @@ mod tests { Err(ModelError::NoLayers) )); } + + #[test] + fn checkpoint_rejects_zero_feature_width() { + let mut map = HashMap::new(); + map.insert("layers.0.weight".into(), Tensor::zeros(vec![2, 0])); + map.insert("layers.0.bias".into(), Tensor::zeros(vec![2])); + assert!(matches!(Mlp::from_tensors(&map), Err(ModelError::BadShape { .. }))); + } } From 3b510a13f7353764b0e0a70f5f3f19f8454c105e Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:48 -0400 Subject: [PATCH 18/24] fix: revalidate public MLP layers before forwarding inputs --- src/model.rs | 38 +++++++++++++++++++++++++++++++++++--- 1 file changed, 35 insertions(+), 3 deletions(-) diff --git a/src/model.rs b/src/model.rs index 52750ee..7dce735 100644 --- a/src/model.rs +++ b/src/model.rs @@ -132,6 +132,29 @@ impl Mlp { Ok(Self { layers }) } + /// Revalidates public layers before inference. + pub fn validate(&self) -> Result<(), ModelError> { + if self.layers.is_empty() { return Err(ModelError::NoLayers); } + let mut previous_width = None; + for (index, (weight, bias)) in self.layers.iter().enumerate() { + weight.validate()?; + bias.validate()?; + if weight.shape.len() != 2 || weight.shape.contains(&0) { + return Err(ModelError::BadShape { name: format!("layers.{index}.weight"), shape: weight.shape.clone() }); + } + if bias.shape.as_slice() != [weight.shape[0]] { + return Err(ModelError::BadShape { name: format!("layers.{index}.bias"), shape: bias.shape.clone() }); + } + if let Some(expected) = previous_width { + if weight.shape[1] != expected { + return Err(ModelError::WidthMismatch { expected, got: weight.shape[1] }); + } + } + previous_width = Some(weight.shape[0]); + } + Ok(()) + } + /// Input feature width expected by the first layer. pub fn in_features(&self) -> usize { self.layers[0].0.shape[1] @@ -151,9 +174,7 @@ impl Mlp { backend: &B, x: &Tensor, ) -> Result { - if self.layers.is_empty() { - return Err(ModelError::NoLayers); - } + self.validate()?; x.validate()?; let (_, in_f) = x.dims2().map_err(|_| ModelError::WidthMismatch { expected: self.in_features(), @@ -272,4 +293,15 @@ mod tests { map.insert("layers.0.bias".into(), Tensor::zeros(vec![2])); assert!(matches!(Mlp::from_tensors(&map), Err(ModelError::BadShape { .. }))); } + + #[test] + fn forward_revalidates_mutated_public_weights_and_biases() { + let mut model = tiny_model(); + let input = Tensor::zeros(vec![1, 2]); + model.layers[0].0.data.pop(); + assert!(matches!(model.forward(&CpuBackend, &input), Err(ModelError::Tensor(_)))); + let mut model = tiny_model(); + model.layers[0].1 = Tensor::zeros(vec![1]); + assert!(matches!(model.forward(&CpuBackend, &input), Err(ModelError::BadShape { .. }))); + } } From 4d6a9f07010ab6f70642f054c64f49518369aa58 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:49 -0400 Subject: [PATCH 19/24] fix: install downloaded GPT-2 assets atomically --- python/fetch_gpt2.py | 18 ++++++++++++++++- python/test_fetch_gpt2.py | 41 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 python/test_fetch_gpt2.py diff --git a/python/fetch_gpt2.py b/python/fetch_gpt2.py index b099448..68e8823 100644 --- a/python/fetch_gpt2.py +++ b/python/fetch_gpt2.py @@ -5,6 +5,7 @@ """ import os +import tempfile import urllib.request BASE = "https://huggingface.co/openai-community/gpt2/resolve/main" @@ -12,6 +13,21 @@ OUT = os.path.join("models", "gpt2") +def download_file(url, destination): + """Install a completed transfer atomically, preserving any previous asset.""" + directory = os.path.dirname(os.path.abspath(destination)) + descriptor, temporary = tempfile.mkstemp(prefix=".gpt2-", suffix=".part", dir=directory) + os.close(descriptor) + try: + urllib.request.urlretrieve(url, temporary) + if os.path.getsize(temporary) == 0: + raise OSError("downloaded file is empty") + os.replace(temporary, destination) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + def main(): os.makedirs(OUT, exist_ok=True) for name in FILES: @@ -20,7 +36,7 @@ def main(): print(f" have {name}") continue print(f" downloading {name} …") - urllib.request.urlretrieve(f"{BASE}/{name}", dst) + download_file(f"{BASE}/{name}", dst) print(f"Done. Weights in {OUT}/") print('Try: cargo run --release --bin generate -- --prompt "Once upon a time"') diff --git a/python/test_fetch_gpt2.py b/python/test_fetch_gpt2.py new file mode 100644 index 0000000..ec2b1d5 --- /dev/null +++ b/python/test_fetch_gpt2.py @@ -0,0 +1,41 @@ +"""Network-free regression tests for GPT-2 asset installation.""" +import os +import tempfile +import unittest +from unittest.mock import patch + +import fetch_gpt2 + + +class DownloadTests(unittest.TestCase): + def test_complete_download_replaces_asset(self): + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + def transfer(url, temporary): + with open(temporary, "wb") as stream: + stream.write(b"complete") + with patch("fetch_gpt2.urllib.request.urlretrieve", side_effect=transfer): + fetch_gpt2.download_file("https://example.test/model", destination) + with open(destination, "rb") as stream: + self.assertEqual(stream.read(), b"complete") + self.assertEqual(os.listdir(directory), ["model"]) + + def test_failure_preserves_previous_asset_and_removes_partial(self): + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with open(destination, "wb") as stream: + stream.write(b"previous") + def transfer(url, temporary): + with open(temporary, "wb") as stream: + stream.write(b"partial") + raise OSError("interrupted") + with patch("fetch_gpt2.urllib.request.urlretrieve", side_effect=transfer): + with self.assertRaises(OSError): + fetch_gpt2.download_file("https://example.test/model", destination) + with open(destination, "rb") as stream: + self.assertEqual(stream.read(), b"previous") + self.assertEqual(os.listdir(directory), ["model"]) + + +if __name__ == "__main__": + unittest.main() From 24fe130b1937412551b76389e34fdb1d318a1aaa Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:51 -0400 Subject: [PATCH 20/24] feat: configure GPT-2 downloads and validate completed transfers --- python/fetch_gpt2.py | 66 ++++++++++++++++++++++++++------------- python/test_fetch_gpt2.py | 55 +++++++++++++++++++++++++------- 2 files changed, 89 insertions(+), 32 deletions(-) diff --git a/python/fetch_gpt2.py b/python/fetch_gpt2.py index 68e8823..116987d 100644 --- a/python/fetch_gpt2.py +++ b/python/fetch_gpt2.py @@ -1,10 +1,9 @@ -"""Download GPT-2 (124M) weights + tokenizer into models/gpt2/. - -Pulls the public HuggingFace `openai-community/gpt2` files (~550 MB) with no -extra dependencies beyond the standard library. These files are gitignored. -""" - +"""Download GPT-2 weights and tokenizer with atomic, bounded transfers.""" +import argparse +import math import os +import shutil +import sys import tempfile import urllib.request @@ -13,33 +12,58 @@ OUT = os.path.join("models", "gpt2") -def download_file(url, destination): +def positive_timeout(value): + try: + seconds = float(value) + except ValueError as error: + raise argparse.ArgumentTypeError("timeout must be a number") from error + if not math.isfinite(seconds) or seconds <= 0: + raise argparse.ArgumentTypeError("timeout must be finite and positive") + return seconds + + +def download_file(url, destination, timeout=60.0): """Install a completed transfer atomically, preserving any previous asset.""" directory = os.path.dirname(os.path.abspath(destination)) descriptor, temporary = tempfile.mkstemp(prefix=".gpt2-", suffix=".part", dir=directory) os.close(descriptor) try: - urllib.request.urlretrieve(url, temporary) - if os.path.getsize(temporary) == 0: + with urllib.request.urlopen(url, timeout=timeout) as response, open(temporary, "wb") as output: + shutil.copyfileobj(response, output) + expected = response.headers.get("Content-Length") + size = os.path.getsize(temporary) + if size == 0: raise OSError("downloaded file is empty") + if expected is not None and size != int(expected): + raise OSError(f"incomplete transfer: expected {expected} bytes, received {size}") os.replace(temporary, destination) finally: if os.path.exists(temporary): os.unlink(temporary) -def main(): - os.makedirs(OUT, exist_ok=True) - for name in FILES: - dst = os.path.join(OUT, name) - if os.path.exists(dst) and os.path.getsize(dst) > 0: - print(f" have {name}") - continue - print(f" downloading {name} …") - download_file(f"{BASE}/{name}", dst) - print(f"Done. Weights in {OUT}/") - print('Try: cargo run --release --bin generate -- --prompt "Once upon a time"') +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", default=OUT, help="asset directory (default: models/gpt2)") + parser.add_argument("--timeout", type=positive_timeout, default=60.0, help="socket timeout in seconds") + parser.add_argument("--force", action="store_true", help="replace existing assets") + args = parser.parse_args(argv) + try: + os.makedirs(args.out, exist_ok=True) + for name in FILES: + destination = os.path.join(args.out, name) + if not args.force and os.path.isfile(destination) and os.path.getsize(destination) > 0: + print(f" have {name}") + continue + print(f" downloading {name} …") + download_file(f"{BASE}/{name}", destination, args.timeout) + except (OSError, ValueError) as error: + print(f"Download failed: {error}", file=sys.stderr) + return 1 + print(f"Done. Weights in {args.out}/") + print(f'Try: cargo run --release --bin generate -- --model-dir "{args.out}" --prompt "Once upon a time"') + return 0 if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/python/test_fetch_gpt2.py b/python/test_fetch_gpt2.py index ec2b1d5..2185a7a 100644 --- a/python/test_fetch_gpt2.py +++ b/python/test_fetch_gpt2.py @@ -1,4 +1,6 @@ """Network-free regression tests for GPT-2 asset installation.""" +import argparse +import io import os import tempfile import unittest @@ -7,35 +9,66 @@ import fetch_gpt2 +class Response(io.BytesIO): + def __init__(self, payload=b"complete", length=None): + super().__init__(payload) + self.headers = {} if length is None else {"Content-Length": str(length)} + + class DownloadTests(unittest.TestCase): def test_complete_download_replaces_asset(self): with tempfile.TemporaryDirectory() as directory: destination = os.path.join(directory, "model") - def transfer(url, temporary): - with open(temporary, "wb") as stream: - stream.write(b"complete") - with patch("fetch_gpt2.urllib.request.urlretrieve", side_effect=transfer): - fetch_gpt2.download_file("https://example.test/model", destination) + with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(length=8)) as opener: + fetch_gpt2.download_file("https://example.test/model", destination, 12.0) + opener.assert_called_once_with("https://example.test/model", timeout=12.0) with open(destination, "rb") as stream: self.assertEqual(stream.read(), b"complete") self.assertEqual(os.listdir(directory), ["model"]) - def test_failure_preserves_previous_asset_and_removes_partial(self): + def test_failure_preserves_previous_asset(self): with tempfile.TemporaryDirectory() as directory: destination = os.path.join(directory, "model") with open(destination, "wb") as stream: stream.write(b"previous") - def transfer(url, temporary): - with open(temporary, "wb") as stream: - stream.write(b"partial") - raise OSError("interrupted") - with patch("fetch_gpt2.urllib.request.urlretrieve", side_effect=transfer): + with patch("fetch_gpt2.urllib.request.urlopen", side_effect=OSError("interrupted")): with self.assertRaises(OSError): fetch_gpt2.download_file("https://example.test/model", destination) with open(destination, "rb") as stream: self.assertEqual(stream.read(), b"previous") self.assertEqual(os.listdir(directory), ["model"]) + def test_empty_and_truncated_downloads_never_become_final_assets(self): + for payload, length in [(b"", None), (b"short", 100)]: + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(payload, length)): + with self.assertRaises(OSError): + fetch_gpt2.download_file("https://example.test/model", destination) + self.assertEqual(os.listdir(directory), []) + + def test_timeout_validation(self): + for value in ["0", "-1", "NaN", "inf", "invalid"]: + with self.assertRaises(argparse.ArgumentTypeError): + fetch_gpt2.positive_timeout(value) + self.assertEqual(fetch_gpt2.positive_timeout("5"), 5.0) + + def test_existing_assets_are_skipped_unless_forced(self): + with tempfile.TemporaryDirectory() as directory: + for name in fetch_gpt2.FILES: + with open(os.path.join(directory, name), "wb") as stream: + stream.write(b"previous") + with patch("builtins.print"), patch("fetch_gpt2.download_file") as transfer: + self.assertEqual(fetch_gpt2.main(["--out", directory]), 0) + transfer.assert_not_called() + self.assertEqual(fetch_gpt2.main(["--out", directory, "--force", "--timeout", "3"]), 0) + self.assertEqual(transfer.call_count, len(fetch_gpt2.FILES)) + + def test_main_reports_network_failure_as_unsuccessful(self): + with tempfile.TemporaryDirectory() as directory: + with patch("builtins.print"), patch("fetch_gpt2.download_file", side_effect=OSError("network")): + self.assertEqual(fetch_gpt2.main(["--out", directory]), 1) + if __name__ == "__main__": unittest.main() From 6bf99872f714517e638f8767facf38d35edb655c Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:43:52 -0400 Subject: [PATCH 21/24] ci: test atomic asset downloads and document numeric contracts --- .github/workflows/ci.yml | 2 ++ docs/correctness.md | 27 +++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 852a4ed..12e5145 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,6 +21,8 @@ jobs: with: components: rustfmt, clippy - uses: Swatinem/rust-cache@v2 + - name: asset downloader tests (no network) + run: python3 -m unittest discover -s python -p 'test_*.py' - name: rustfmt run: cargo fmt --all -- --check - name: clippy diff --git a/docs/correctness.md b/docs/correctness.md index 3adf6aa..4ce17e1 100644 --- a/docs/correctness.md +++ b/docs/correctness.md @@ -93,3 +93,30 @@ exist: - No SSM `scan` op (not implemented). - No full transformer-LM end-to-end parity (no LM yet) — only the MLP is verified end-to-end, plus every individual operator above. + +## Invalid-input contracts and sampling + +CPU reference operators check dimension products before allocation. Normalization +requires a positive width and a finite non-negative epsilon; rotary embeddings +require a finite positive theta. Attention validates head widths and bounds +absolute causal offsets without wrapping. + +Use Sampler::sample_token for checked single-token sampling, and Sampler::validate +to check temperature. Empty or non-finite logits are rejected. Probabilities use +f64 intermediates; top-k ties prefer smaller token IDs and the retained candidates +have a stable ID order. Seeded sampled sequences may differ from earlier versions +because their candidate ordering and arithmetic were corrected. Greedy finite +argmax retains its earliest-ID tie behavior. + +Tokenizer merge parsing retains hash-prefixed merge tokens, rejects malformed or +duplicate pairs, and checks both operands. Repeated whitespace preserves GPT-2's +lookahead boundary before the next token. The additive try_decode and +try_decode_bytes APIs report unknown token IDs; existing decode APIs retain their +permissive behavior. + +The download script accepts --out, --timeout, and --force. Each transfer uses a +temporary file beside its destination, checks any advertised Content-Length, and +atomically replaces the destination only after completion. Existing nonempty +assets are skipped by default; --force refreshes them. Network-free downloader +regressions run in CI. These tests do not download or validate pretrained GPT-2 +weights and do not replace on-device Metal parity tests. From 65d801398646a6332fce2dd57bb0300f0fc759b1 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:46:11 -0400 Subject: [PATCH 22/24] style: format numeric and tokenizer regressions --- src/gpt2.rs | 102 ++++++++++++++++++++++++++++++++++++++--------- src/model.rs | 34 ++++++++++++---- src/ops.rs | 37 +++++++++++++---- src/tokenizer.rs | 40 +++++++++++++++---- 4 files changed, 172 insertions(+), 41 deletions(-) diff --git a/src/gpt2.rs b/src/gpt2.rs index e74e715..2253d70 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -368,12 +368,20 @@ impl Gpt2 { mut on_token: impl FnMut(usize) -> bool, ) -> Vec { let mut toks = prompt.to_vec(); - if max_new == 0 { return toks; } + if max_new == 0 { + return toks; + } self.config.validate().expect("invalid GPT-2 configuration"); sampler.validate().expect("invalid generation sampler"); assert!(!prompt.is_empty(), "generation requires a non-empty prompt"); - assert!(eot_token < self.config.vocab_size, "end-of-text ID exceeds GPT-2 vocabulary"); - assert!(prompt.iter().all(|&token| token < self.config.vocab_size), "prompt ID exceeds GPT-2 vocabulary"); + assert!( + eot_token < self.config.vocab_size, + "end-of-text ID exceeds GPT-2 vocabulary" + ); + assert!( + prompt.iter().all(|&token| token < self.config.vocab_size), + "prompt ID exceeds GPT-2 vocabulary" + ); let n_ctx = self.config.n_ctx; let mut rng = Rng::new(sampler.seed); for _ in 0..max_new { @@ -428,12 +436,18 @@ pub struct Sampler { impl Sampler { /// Greedy (argmax) sampling. pub fn greedy() -> Self { - Sampler { temperature: 0.0, top_k: 0, seed: 0 } + Sampler { + temperature: 0.0, + top_k: 0, + seed: 0, + } } pub fn validate(&self) -> Result<(), Gpt2Error> { if !self.temperature.is_finite() || self.temperature < 0.0 { - return Err(Gpt2Error::InvalidSampling("temperature must be finite and non-negative")); + return Err(Gpt2Error::InvalidSampling( + "temperature must be finite and non-negative", + )); } Ok(()) } @@ -444,7 +458,8 @@ impl Sampler { } fn sample(&self, logits: &[f32], rng: &mut Rng) -> usize { - self.sample_checked(logits, rng).expect("invalid token sampling input") + self.sample_checked(logits, rng) + .expect("invalid token sampling input") } fn sample_checked(&self, logits: &[f32], rng: &mut Rng) -> Result { @@ -466,10 +481,18 @@ impl Sampler { idx.truncate(self.top_k); idx.sort_unstable(); } - let max = idx.iter().map(|&i| logits[i]).fold(f32::NEG_INFINITY, f32::max); - let mut probs: Vec = idx.iter().map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()).collect(); + let max = idx + .iter() + .map(|&i| logits[i]) + .fold(f32::NEG_INFINITY, f32::max); + let mut probs: Vec = idx + .iter() + .map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()) + .collect(); let sum: f64 = probs.iter().sum(); - for p in &mut probs { *p /= sum; } + for p in &mut probs { + *p /= sum; + } Ok(select_probability(&idx, &probs, f64::from(rng.next_f32()))) } } @@ -478,10 +501,15 @@ fn select_probability(indices: &[usize], probabilities: &[f64], draw: f64) -> us let mut cumulative = 0.0; for (index, &probability) in probabilities.iter().enumerate() { cumulative += probability; - if draw < cumulative { return indices[index]; } + if draw < cumulative { + return indices[index]; + } } // Rounding near one must still select a candidate with positive weight. - let last = probabilities.iter().rposition(|&p| p > 0.0).expect("sampling has no positive weight"); + let last = probabilities + .iter() + .rposition(|&p| p > 0.0) + .expect("sampling has no positive weight"); indices[last] } @@ -617,7 +645,13 @@ mod tests { #[test] fn checked_sampling_rejects_invalid_inputs() { for temperature in [f32::NAN, f32::INFINITY, -1.0] { - assert!(Sampler { temperature, top_k: 1, seed: 0 }.sample_token(&[0.0]).is_err()); + assert!(Sampler { + temperature, + top_k: 1, + seed: 0 + } + .sample_token(&[0.0]) + .is_err()); } assert!(Sampler::greedy().sample_token(&[]).is_err()); for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { @@ -627,15 +661,26 @@ mod tests { #[test] fn extreme_finite_logits_sample_without_overflow() { - let sampler = Sampler { temperature: f32::MIN_POSITIVE, top_k: 0, seed: 42 }; + let sampler = Sampler { + temperature: f32::MIN_POSITIVE, + top_k: 0, + seed: 42, + }; assert_eq!(sampler.sample_token(&[-f32::MAX, f32::MAX]).unwrap(), 1); } #[test] fn top_k_ties_choose_lowest_ids_and_repeat_from_a_seed() { - let sampler = Sampler { temperature: 1.0, top_k: 1, seed: 42 }; + let sampler = Sampler { + temperature: 1.0, + top_k: 1, + seed: 42, + }; assert_eq!(sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(), 0); - let sampler = Sampler { top_k: 2, ..sampler }; + let sampler = Sampler { + top_k: 2, + ..sampler + }; let first = sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(); assert!(first < 2); assert_eq!(first, sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap()); @@ -644,18 +689,39 @@ mod tests { #[test] fn sampling_boundaries_never_choose_zero_weight() { assert_eq!(select_probability(&[10, 20], &[0.0, 1.0], 0.0), 20); - assert_eq!(select_probability(&[10, 20, 30], &[0.5, 0.4999999999999999, 0.0], 0.9999999999999999), 20); + assert_eq!( + select_probability( + &[10, 20, 30], + &[0.5, 0.4999999999999999, 0.0], + 0.9999999999999999 + ), + 20 + ); } #[test] fn zero_token_generation_returns_the_prompt_without_inference() { - let tokens = forward_test_model().generate(&crate::model::CpuBackend, &[], 0, &Sampler::greedy(), 1, |_| panic!("unexpected inference")); + let tokens = forward_test_model().generate( + &crate::model::CpuBackend, + &[], + 0, + &Sampler::greedy(), + 1, + |_| panic!("unexpected inference"), + ); assert!(tokens.is_empty()); } #[test] #[should_panic(expected = "end-of-text ID exceeds")] fn generation_rejects_unknown_stop_tokens() { - forward_test_model().generate(&crate::model::CpuBackend, &[0], 1, &Sampler::greedy(), 999, |_| {}); + forward_test_model().generate( + &crate::model::CpuBackend, + &[0], + 1, + &Sampler::greedy(), + 999, + |_| {}, + ); } } diff --git a/src/model.rs b/src/model.rs index 7dce735..de24a4a 100644 --- a/src/model.rs +++ b/src/model.rs @@ -134,20 +134,31 @@ impl Mlp { /// Revalidates public layers before inference. pub fn validate(&self) -> Result<(), ModelError> { - if self.layers.is_empty() { return Err(ModelError::NoLayers); } + if self.layers.is_empty() { + return Err(ModelError::NoLayers); + } let mut previous_width = None; for (index, (weight, bias)) in self.layers.iter().enumerate() { weight.validate()?; bias.validate()?; if weight.shape.len() != 2 || weight.shape.contains(&0) { - return Err(ModelError::BadShape { name: format!("layers.{index}.weight"), shape: weight.shape.clone() }); + return Err(ModelError::BadShape { + name: format!("layers.{index}.weight"), + shape: weight.shape.clone(), + }); } if bias.shape.as_slice() != [weight.shape[0]] { - return Err(ModelError::BadShape { name: format!("layers.{index}.bias"), shape: bias.shape.clone() }); + return Err(ModelError::BadShape { + name: format!("layers.{index}.bias"), + shape: bias.shape.clone(), + }); } if let Some(expected) = previous_width { if weight.shape[1] != expected { - return Err(ModelError::WidthMismatch { expected, got: weight.shape[1] }); + return Err(ModelError::WidthMismatch { + expected, + got: weight.shape[1], + }); } } previous_width = Some(weight.shape[0]); @@ -291,7 +302,10 @@ mod tests { let mut map = HashMap::new(); map.insert("layers.0.weight".into(), Tensor::zeros(vec![2, 0])); map.insert("layers.0.bias".into(), Tensor::zeros(vec![2])); - assert!(matches!(Mlp::from_tensors(&map), Err(ModelError::BadShape { .. }))); + assert!(matches!( + Mlp::from_tensors(&map), + Err(ModelError::BadShape { .. }) + )); } #[test] @@ -299,9 +313,15 @@ mod tests { let mut model = tiny_model(); let input = Tensor::zeros(vec![1, 2]); model.layers[0].0.data.pop(); - assert!(matches!(model.forward(&CpuBackend, &input), Err(ModelError::Tensor(_)))); + assert!(matches!( + model.forward(&CpuBackend, &input), + Err(ModelError::Tensor(_)) + )); let mut model = tiny_model(); model.layers[0].1 = Tensor::zeros(vec![1]); - assert!(matches!(model.forward(&CpuBackend, &input), Err(ModelError::BadShape { .. }))); + assert!(matches!( + model.forward(&CpuBackend, &input), + Err(ModelError::BadShape { .. }) + )); } } diff --git a/src/ops.rs b/src/ops.rs index fbd690d..e7dda01 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -11,7 +11,8 @@ //! * `rope` uses the rotate-half (GPT-NeoX / HF) convention. fn checked_elements(rows: usize, cols: usize) -> usize { - rows.checked_mul(cols).expect("operator dimensions overflow") + rows.checked_mul(cols) + .expect("operator dimensions overflow") } /// Standard matrix multiply: `A`[m×k] · `B`[k×n] → `C`[m×n] (row-major). @@ -100,7 +101,10 @@ pub fn layernorm( eps: f32, ) -> Vec { assert!(d > 0, "normalization width must be positive"); - assert!(eps.is_finite() && eps >= 0.0, "normalization epsilon must be finite and non-negative"); + assert!( + eps.is_finite() && eps >= 0.0, + "normalization epsilon must be finite and non-negative" + ); assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(gamma.len(), d); assert_eq!(beta.len(), d); @@ -120,7 +124,10 @@ pub fn layernorm( /// Row-wise RMSNorm over the last dimension of size `d`: `y = x / sqrt(mean(x²) + eps) * gamma`. pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec { assert!(d > 0, "normalization width must be positive"); - assert!(eps.is_finite() && eps >= 0.0, "normalization epsilon must be finite and non-negative"); + assert!( + eps.is_finite() && eps >= 0.0, + "normalization epsilon must be finite and non-negative" + ); assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(gamma.len(), d); let mut out = vec![0.0f32; checked_elements(rows, d)]; @@ -140,7 +147,10 @@ pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec /// `d` must be even. pub fn rope_inplace(x: &mut [f32], positions: &[usize], rows: usize, d: usize, theta: f32) { assert!(d > 0, "rope head dim must be positive"); - assert!(theta.is_finite() && theta > 0.0, "rope theta must be finite and positive"); + assert!( + theta.is_finite() && theta > 0.0, + "rope theta must be finite and positive" + ); assert_eq!(x.len(), checked_elements(rows, d)); assert_eq!(positions.len(), rows); assert_eq!(d % 2, 0, "rope head dim must be even"); @@ -177,7 +187,10 @@ pub fn attention( q_offset: usize, ) -> Vec { assert!(d > 0, "attention width must be positive"); - assert!(seq > 0 || m == 0, "non-empty attention queries require keys"); + assert!( + seq > 0 || m == 0, + "non-empty attention queries require keys" + ); assert_eq!(q.len(), checked_elements(m, d)); assert_eq!(k.len(), checked_elements(seq, d)); assert_eq!(v.len(), checked_elements(seq, d)); @@ -213,7 +226,10 @@ pub fn attention( /// heads laid out contiguously per row; output has the same shape. Query `i` /// attends causally over keys `j <= i`. This is the GPT-2 attention reference. pub fn mha(q: &[f32], k: &[f32], v: &[f32], seq: usize, heads: usize, head_dim: usize) -> Vec { - assert!(heads > 0 && head_dim > 0, "attention heads and head width must be positive"); + assert!( + heads > 0 && head_dim > 0, + "attention heads and head width must be positive" + ); let hd = checked_elements(heads, head_dim); assert_eq!(q.len(), checked_elements(seq, hd)); assert_eq!(k.len(), checked_elements(seq, hd)); @@ -392,7 +408,9 @@ mod tests { #[test] fn normalization_rejects_invalid_epsilon() { for eps in [f32::NAN, f32::INFINITY, -1.0] { - assert!(std::panic::catch_unwind(|| layernorm(&[1.0], &[1.0], &[0.0], 1, 1, eps)).is_err()); + assert!( + std::panic::catch_unwind(|| layernorm(&[1.0], &[1.0], &[0.0], 1, 1, eps)).is_err() + ); assert!(std::panic::catch_unwind(|| rmsnorm(&[1.0], &[1.0], 1, 1, eps)).is_err()); } } @@ -407,7 +425,10 @@ mod tests { #[test] fn rope_rejects_invalid_frequency_bases() { for theta in [0.0, -1.0, f32::NAN, f32::INFINITY] { - assert!(std::panic::catch_unwind(|| rope_inplace(&mut [1.0, 2.0], &[0], 1, 2, theta)).is_err()); + assert!( + std::panic::catch_unwind(|| rope_inplace(&mut [1.0, 2.0], &[0], 1, 2, theta)) + .is_err() + ); } } diff --git a/src/tokenizer.rs b/src/tokenizer.rs index a051fa2..56e28c0 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -117,7 +117,10 @@ impl Tokenizer { if !vocab.contains_key(&merged) { return Err(TokenizerError::MissingVocabSymbol(merged)); } - if bpe_ranks.insert((a.to_string(), b.to_string()), rank).is_some() { + if bpe_ranks + .insert((a.to_string(), b.to_string()), rank) + .is_some() + { return Err(TokenizerError::DuplicateMerge(line.to_string())); } } else { @@ -171,8 +174,15 @@ impl Tokenizer { pub fn try_decode_bytes(&self, ids: &[usize]) -> Result, TokenizerError> { let mut bytes = Vec::new(); for id in ids { - let token = self.decoder.get(id).ok_or(TokenizerError::UnknownTokenId(*id))?; - bytes.extend(token.chars().filter_map(|c| self.byte_decoder.get(&c).copied())); + let token = self + .decoder + .get(id) + .ok_or(TokenizerError::UnknownTokenId(*id))?; + bytes.extend( + token + .chars() + .filter_map(|c| self.byte_decoder.get(&c).copied()), + ); } Ok(bytes) } @@ -419,7 +429,10 @@ mod tests { #[test] fn malformed_merge_lines_are_reported() { for line in ["a", "a b c"] { - assert!(matches!(Tokenizer::from_assets(byte_vocab(), line), Err(TokenizerError::InvalidMergeLine(_)))); + assert!(matches!( + Tokenizer::from_assets(byte_vocab(), line), + Err(TokenizerError::InvalidMergeLine(_)) + )); } } @@ -427,21 +440,32 @@ mod tests { fn duplicate_merges_are_rejected() { let mut vocab = byte_vocab(); vocab.insert("ab".into(), 256); - assert!(matches!(Tokenizer::from_assets(vocab, "a b\na b\n"), Err(TokenizerError::DuplicateMerge(_)))); + assert!(matches!( + Tokenizer::from_assets(vocab, "a b\na b\n"), + Err(TokenizerError::DuplicateMerge(_)) + )); } #[test] fn missing_merge_operands_are_reported() { let mut vocab = byte_vocab(); vocab.insert("not_presenta".into(), 256); - assert!(matches!(Tokenizer::from_assets(vocab, "not_present a"), Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "not_present")); + assert!( + matches!(Tokenizer::from_assets(vocab, "not_present a"), Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "not_present") + ); } #[test] fn checked_decode_reports_unknown_ids_and_preserves_bytes() { let tokenizer = Tokenizer::from_assets(byte_vocab(), "").unwrap(); - assert!(matches!(tokenizer.try_decode(&[999]), Err(TokenizerError::UnknownTokenId(999)))); - assert_eq!(tokenizer.try_decode_bytes(&[0xe2, 0x82, 0xac]).unwrap(), vec![0xe2, 0x82, 0xac]); + assert!(matches!( + tokenizer.try_decode(&[999]), + Err(TokenizerError::UnknownTokenId(999)) + )); + assert_eq!( + tokenizer.try_decode_bytes(&[0xe2, 0x82, 0xac]).unwrap(), + vec![0xe2, 0x82, 0xac] + ); assert_eq!(tokenizer.try_decode(&[0xe2, 0x82, 0xac]).unwrap(), "€"); } } From 3afe197d285593886b717a3a1d13fa17dee0cd73 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:47:49 -0400 Subject: [PATCH 23/24] fix: treat signed zero as a sampling tie --- src/gpt2.rs | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/gpt2.rs b/src/gpt2.rs index 2253d70..cb2f743 100644 --- a/src/gpt2.rs +++ b/src/gpt2.rs @@ -476,7 +476,11 @@ impl Sampler { let mut idx: Vec = (0..logits.len()).collect(); if self.top_k > 0 && self.top_k < logits.len() { idx.select_nth_unstable_by(self.top_k, |&a, &b| { - logits[b].total_cmp(&logits[a]).then_with(|| a.cmp(&b)) + if logits[a] == logits[b] { + a.cmp(&b) + } else { + logits[b].total_cmp(&logits[a]) + } }); idx.truncate(self.top_k); idx.sort_unstable(); @@ -724,4 +728,14 @@ mod tests { |_| {}, ); } + + #[test] + fn top_k_ties_treat_signed_zero_as_equal() { + let sampler = Sampler { + temperature: 1.0, + top_k: 1, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[-0.0, 0.0]).unwrap(), 0); + } } From 376cf34feb6a69947fb16ed1d67c7218f01479a6 Mon Sep 17 00:00:00 2001 From: realyashnegi <54710562+yash27-lab@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:47:50 -0400 Subject: [PATCH 24/24] perf: return empty matrix outputs without iterating rows --- src/ops.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/ops.rs b/src/ops.rs index e7dda01..cd4907c 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -20,6 +20,9 @@ pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { assert_eq!(a.len(), checked_elements(m, k), "A has wrong length"); assert_eq!(b.len(), checked_elements(k, n), "B has wrong length"); let mut c = vec![0.0f32; checked_elements(m, n)]; + if c.is_empty() { + return c; + } for row in 0..m { for i in 0..k { let a_ik = a[row * k + i]; @@ -43,6 +46,9 @@ pub fn linear(x: &[f32], w: &[f32], b: &[f32], n: usize, in_f: usize, out_f: usi assert_eq!(w.len(), checked_elements(out_f, in_f)); assert_eq!(b.len(), out_f); let mut y = vec![0.0f32; checked_elements(n, out_f)]; + if y.is_empty() { + return y; + } for row in 0..n { for o in 0..out_f { let mut acc = b[o]; @@ -266,6 +272,9 @@ pub fn dequantize_int8(q: &[i8], scales: &[f32], rows: usize, cols: usize) -> Ve assert_eq!(q.len(), checked_elements(rows, cols)); assert_eq!(scales.len(), rows); let mut out = vec![0.0f32; checked_elements(rows, cols)]; + if out.is_empty() { + return out; + } for r in 0..rows { let s = scales[r]; for c in 0..cols { @@ -444,4 +453,11 @@ mod tests { assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, usize::MAX, 2)).is_err()); assert!(mha(&[], &[], &[], 0, 1, 1).is_empty()); } + + #[test] + fn empty_matrix_outputs_skip_unused_rows() { + assert!(matmul(&[], &[], usize::MAX, 0, 0).is_empty()); + assert!(linear(&[], &[], &[], usize::MAX, 0, 0).is_empty()); + assert!(dequantize_int8(&[], &[1.0; 3], 3, 0).is_empty()); + } }