foxhunt/crates/ml-supervised/src/kan/network.rs

//! KAN network: a stack of KAN layers.
//!
//! Chains multiple KAN layers according to `layer_widths` from the config.
//! No explicit activation between layers because each KAN layer contains
//! its own learned activation functions (B-splines).

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::GpuTensor;

use super::config::KANConfig;
use super::layer::KANLayer;

/// A multi-layer Kolmogorov-Arnold Network.
#[derive(Debug)]
pub struct KANNetwork {
    layers: Vec<KANLayer>,
}

impl KANNetwork {
    /// Build a KAN network from config.
    ///
    /// Creates one KAN layer per consecutive pair of widths in `config.layer_widths`.
    pub fn new(config: &KANConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if config.layer_widths.len() < 2 {
            return Err(MLError::ConfigError(
                "layer_widths must have at least 2 entries".to_owned(),
            ));
        }
        if config.layer_widths.contains(&0) {
            return Err(MLError::ConfigError(
                "KAN requires all layer_widths > 0".to_owned(),
            ));
        }

        let mut layers = Vec::with_capacity(config.layer_widths.len() - 1);
        for i in 0..config.layer_widths.len() - 1 {
            let in_dim = *config.layer_widths.get(i).ok_or_else(|| {
                MLError::ConfigError(format!("Missing layer_widths index {}", i))
            })?;
            let out_dim = *config.layer_widths.get(i + 1).ok_or_else(|| {
                MLError::ConfigError(format!("Missing layer_widths index {}", i + 1))
            })?;

            let layer = KANLayer::new(
                in_dim,
                out_dim,
                config.grid_size,
                config.spline_order,
                stream,
            )?;
            layers.push(layer);
        }

        Ok(Self { layers })
    }

    /// Forward pass: chains all layers sequentially.
    ///
    /// Input shape: `(batch, layer_widths[0])`
    /// Output shape: `(batch, layer_widths[last])`
    pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor, MLError> {
        let mut x = input.clone();
        for layer in &self.layers {
            x = layer.forward(&x)?;
        }
        Ok(x)
    }

    /// Access the layers for parameter extraction (checkpoint/weight access).
    pub fn layers(&self) -> &[KANLayer] {
        &self.layers
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA required");
        ctx.new_stream().expect("Failed to create stream")
    }

    fn make_config() -> KANConfig {
        KANConfig {
            grid_size: 3,
            spline_order: 3,
            layer_widths: vec![8, 4, 1],
            learning_rate: 1e-3,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }

    #[test]
    fn test_kan_network_forward() {
        let stream = test_stream();
        let config = make_config();
        let net = KANNetwork::new(&config, &stream).unwrap();

        let input = GpuTensor::randn(&[2, 8], 0.5, &stream).unwrap();
        let output = net.forward(&input).unwrap();
        assert_eq!(output.shape, vec![2, 1]);
    }

    #[test]
    fn test_kan_network_finite_outputs() {
        let stream = test_stream();
        let config = make_config();
        let net = KANNetwork::new(&config, &stream).unwrap();

        let input = GpuTensor::randn(&[4, 8], 0.5, &stream).unwrap();
        let output = net.forward(&input).unwrap();
        let v = output.to_vec().unwrap(); // test-only readback
        for val in &v {
            assert!(val.is_finite(), "Non-finite output: {}", val);
        }
    }
}