gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(ml): replace unwrap() with ok_or_else() in TFT quantized attention forward pass

Browse Source 
Replace 10 .as_ref().unwrap() calls on Option<QuantizedTensor> weight fields
(q_weights, k_weights, v_weights, o_weights) and attention_cache with
.as_ref().ok_or_else(|| MLError::ModelError(...))? to satisfy the
clippy::unwrap_used deny lint. Affected methods: build_cache,
forward_with_mask, compute_projections_slow, and the test
test_attention_weights_sum_to_one.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-21 23:34:17 +01:00
parent aa130e9554
commit a9ddff03b6
1 changed files with 33 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
ml/src/tft/quantized_attention.rs
View File

@@ -139,16 +139,24 @@ impl QuantizedTemporalAttention {
// Dequantize all weights
let q_weight = self
.quantizer
.dequantize_tensor(self.q_weights.as_ref().unwrap())?;
.dequantize_tensor(self.q_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized Q weights not initialized".to_string())
})?)?;
let k_weight = self
.quantizer
.dequantize_tensor(self.k_weights.as_ref().unwrap())?;
.dequantize_tensor(self.k_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized K weights not initialized".to_string())
})?)?;
let v_weight = self
.quantizer
.dequantize_tensor(self.v_weights.as_ref().unwrap())?;
.dequantize_tensor(self.v_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized V weights not initialized".to_string())
})?)?;
let o_weight = self
.quantizer
.dequantize_tensor(self.o_weights.as_ref().unwrap())?;
.dequantize_tensor(self.o_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized O weights not initialized".to_string())
})?)?;
// Store in cache
self.attention_cache = Some(AttentionWeightCache {
@@ -229,7 +237,9 @@ impl QuantizedTemporalAttention {
// In production, cache should be built once after weight initialization
self.compute_projections_slow(x)?
} else {
let cache = self.attention_cache.as_ref().unwrap();
let cache = self.attention_cache.as_ref().ok_or_else(|| {
MLError::ModelError("attention weight cache not built".to_string())
})?;
// Reshape 3D input to 2D for batch matmul
let dims = x.dims();
@@ -312,7 +322,9 @@ impl QuantizedTemporalAttention {
// Step 8: Output projection
let output = if self.cache_enabled && self.attention_cache.is_some() {
let cache = self.attention_cache.as_ref().unwrap();
let cache = self.attention_cache.as_ref().ok_or_else(|| {
MLError::ModelError("attention weight cache not built".to_string())
})?;
// Reshape attended to 2D for matmul, then reshape back
let attended_2d = attended.reshape(&[batch_size * seq_len, self.hidden_dim])?;
let output_2d = attended_2d.matmul(&cache.o_weight.t()?)?;
@@ -320,7 +332,9 @@ impl QuantizedTemporalAttention {
} else {
let o_weight = self
.quantizer
.dequantize_tensor(self.o_weights.as_ref().unwrap())?;
.dequantize_tensor(self.o_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized O weights not initialized".to_string())
})?)?;
// Reshape attended to 2D for matmul, then reshape back
let attended_2d = attended.reshape(&[batch_size * seq_len, self.hidden_dim])?;
let output_2d = attended_2d.matmul(&o_weight.t()?)?;
@@ -334,13 +348,19 @@ impl QuantizedTemporalAttention {
fn compute_projections_slow(&self, x: &Tensor) -> Result<(Tensor, Tensor, Tensor), MLError> {
let q_weight = self
.quantizer
.dequantize_tensor(self.q_weights.as_ref().unwrap())?;
.dequantize_tensor(self.q_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized Q weights not initialized".to_string())
})?)?;
let k_weight = self
.quantizer
.dequantize_tensor(self.k_weights.as_ref().unwrap())?;
.dequantize_tensor(self.k_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized K weights not initialized".to_string())
})?)?;
let v_weight = self
.quantizer
.dequantize_tensor(self.v_weights.as_ref().unwrap())?;
.dequantize_tensor(self.v_weights.as_ref().ok_or_else(|| {
MLError::ModelError("quantized V weights not initialized".to_string())
})?)?;
// Reshape 3D input to 2D for batch matmul: [batch, seq_len, hidden_dim] -> [batch * seq_len, hidden_dim]
// Candle doesn't support direct 3D x 2D matmul, so we flatten the batch and sequence dimensions
@@ -466,7 +486,9 @@ mod tests {
attention.build_cache()?;
// Manually verify attention weights sum to 1.0
let cache = attention.attention_cache.as_ref().unwrap();
let cache = attention.attention_cache.as_ref().ok_or_else(|| {
MLError::ModelError("attention weight cache not built".to_string())
})?;
let num_heads = 8;
let head_dim = hidden_dim / num_heads;