feat(liquid): add CandleCfCNetwork with full sequence processing
Wraps CfCCell with output projection for end-to-end sequence processing. forward_sequence unrolls the cell over [batch, seq_len, features] input and projects the final hidden state to [batch, output_size]. Includes forward() convenience method with default dt=0.01 for UnifiedTrainable compatibility. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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<Self, MLError> {
|
||||
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<Tensor, MLError> {
|
||||
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<Tensor, MLError> {
|
||||
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]);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user