diff --git a/ml/src/liquid/candle_cfc.rs b/ml/src/liquid/candle_cfc.rs index 20afc7e62..32b86490e 100644 --- a/ml/src/liquid/candle_cfc.rs +++ b/ml/src/liquid/candle_cfc.rs @@ -237,6 +237,107 @@ impl CfCCell { } } +/// Complete CfC v2 Network for sequence processing +/// +/// Wraps `CfCCell` with an output projection layer. Given a 3D input tensor +/// `[batch, seq_len, features]`, unrolls the cell over time steps and projects +/// the final hidden state to `[batch, output_size]`. +#[allow(missing_debug_implementations)] +pub struct CandleCfCNetwork { + cell: CfCCell, + output_layer: Linear, + config: CfCTrainConfig, +} + +impl CandleCfCNetwork { + pub fn new(config: &CfCTrainConfig, vb: &VarBuilder<'_>) -> Result { + let cell = CfCCell::new(config, &vb.pp("cfc"))?; + let output_layer = + candle_nn::linear(config.hidden_size, config.output_size, vb.pp("output")) + .map_err(|e| MLError::ModelError(format!("Output layer init: {}", e)))?; + Ok(Self { + cell, + output_layer, + config: config.clone(), + }) + } + + /// Process full sequence `[batch, seq_len, features]` -> `[batch, output_size]` + pub fn forward_sequence(&self, input: &Tensor, dt: f32) -> Result { + let dims = input.dims(); + if dims.len() != 3 { + return Err(MLError::InvalidInput(format!( + "Expected 3D input [batch, seq_len, features], got {:?}", + dims + ))); + } + let batch_size = *dims.first().ok_or_else(|| { + MLError::InvalidInput("Input tensor has no dimensions".to_string()) + })?; + let seq_len = *dims.get(1).ok_or_else(|| { + MLError::InvalidInput("Input tensor missing seq_len dimension".to_string()) + })?; + let device = input.device(); + + let mut h = Tensor::zeros( + (batch_size, self.config.hidden_size), + DType::F32, + device, + ) + .map_err(|e| MLError::InferenceError(format!("CfC init hidden: {}", e)))?; + + for t in 0..seq_len { + let x_t = input + .narrow(1, t, 1) + .map_err(|e| MLError::InferenceError(format!("CfC narrow t={}: {}", t, e)))? + .squeeze(1) + .map_err(|e| MLError::InferenceError(format!("CfC squeeze t={}: {}", t, e)))?; + h = self.cell.step(&x_t, &h, dt)?; + } + + let output = self + .output_layer + .forward(&h) + .map_err(|e| MLError::InferenceError(format!("CfC output: {}", e)))?; + Ok(output) + } + + /// Forward compatible with UnifiedTrainable (3D input, default dt=0.01) + pub fn forward(&self, input: &Tensor) -> Result { + self.forward_sequence(input, 0.01) + } + + /// Approximate parameter count for reporting purposes + pub fn param_count(&self) -> usize { + let backbone_params = + self.config + .backbone_hidden_sizes + .iter() + .enumerate() + .fold(0, |acc, (i, &size)| { + let in_d = if i == 0 { + self.config.input_size + self.config.hidden_size + } else { + self.config + .backbone_hidden_sizes + .get(i.saturating_sub(1)) + .copied() + .unwrap_or(size) + }; + acc + in_d * size + size + }); + let last_h = self + .config + .backbone_hidden_sizes + .last() + .copied() + .unwrap_or(0); + let heads = 2 * (last_h * last_h + last_h); + let output = self.config.hidden_size * self.config.output_size + self.config.output_size; + backbone_params + heads + output + } +} + #[cfg(test)] mod tests { use super::*; @@ -431,4 +532,57 @@ mod tests { assert!(h1_norm.is_finite()); assert!(h2_norm.is_finite()); } + + // --------------------------------------------------------------- + // Task 4: CandleCfCNetwork tests + // --------------------------------------------------------------- + + #[test] + fn test_cfc_network_creation() { + let device = Device::Cpu; + let varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let config = CfCTrainConfig::default(); + let network = CandleCfCNetwork::new(&config, &vb).unwrap(); + assert!(network.param_count() > 0); + } + + #[test] + fn test_cfc_network_forward_sequence() { + let device = Device::Cpu; + let varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let config = CfCTrainConfig { + input_size: 16, + hidden_size: 32, + output_size: 3, + backbone_hidden_sizes: vec![32], + seq_len: 10, + batch_size: 4, + ..CfCTrainConfig::default() + }; + let network = CandleCfCNetwork::new(&config, &vb).unwrap(); + let input = Tensor::randn(0f32, 1.0, (4, 10, 16), &device).unwrap(); + let output = network.forward_sequence(&input, 0.01).unwrap(); + assert_eq!(output.dims(), &[4, 3]); + } + + #[test] + fn test_cfc_network_forward_trait() { + let device = Device::Cpu; + let varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let config = CfCTrainConfig { + input_size: 16, + hidden_size: 32, + output_size: 3, + backbone_hidden_sizes: vec![32], + seq_len: 10, + ..CfCTrainConfig::default() + }; + let network = CandleCfCNetwork::new(&config, &vb).unwrap(); + let input = Tensor::randn(0f32, 1.0, (4, 10, 16), &device).unwrap(); + let output = network.forward(&input).unwrap(); + assert_eq!(output.dims(), &[4, 3]); + } }