From c88e53c3b3de951ce4ac8fcd47ac8bf572caa8ca Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 21 Feb 2026 21:44:38 +0100 Subject: [PATCH] feat(ml): implement TFT layer normalization in quantized GRN Replace the no-op apply_layer_norm() that returned input unchanged with real layer normalization using layer_norm_with_fallback. Initialize LayerNormParams with proper weight (ones) and bias (zeros) tensors in from_grn constructor instead of None placeholders. --- ml/src/tft/quantized_grn.rs | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/ml/src/tft/quantized_grn.rs b/ml/src/tft/quantized_grn.rs index 55262a116..21f4ac36f 100644 --- a/ml/src/tft/quantized_grn.rs +++ b/ml/src/tft/quantized_grn.rs @@ -92,10 +92,15 @@ impl QuantizedGatedResidualNetwork { }; // Store layer norm parameters (keep in F32) + // Initialize weight=ones and bias=zeros for proper layer normalization + let ln_weight = Tensor::ones(grn.output_dim, candle_core::DType::F32, &device) + .map_err(|e| MLError::ModelError(format!("Failed to create ln_weight: {}", e)))?; + let ln_bias = Tensor::zeros(grn.output_dim, candle_core::DType::F32, &device) + .map_err(|e| MLError::ModelError(format!("Failed to create ln_bias: {}", e)))?; let layer_norm = Some(LayerNormParams { normalized_shape: vec![grn.output_dim], - weight: None, // Would extract from grn.layer_norm in production - bias: None, + weight: Some(ln_weight), + bias: Some(ln_bias), eps: 1e-5, }); @@ -221,12 +226,17 @@ impl QuantizedGatedResidualNetwork { /// Apply layer normalization (kept in F32) fn apply_layer_norm(&self, x: &Tensor) -> Result { - // Simplified layer norm for TDD - // In production, would use actual layer_norm with weights/bias - - // For now, return input (layer norm would normalize here) - // TODO: Implement proper layer normalization - Ok(x.clone()) + if let Some(ln) = &self.layer_norm { + crate::cuda_compat::layer_norm_with_fallback( + x, + &ln.normalized_shape, + ln.weight.as_ref(), + ln.bias.as_ref(), + ln.eps, + ) + } else { + Ok(x.clone()) + } } /// Get quantization type