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

//! KAN layer implementation.
//!
//! A single KAN layer maps `(batch, in_dim)` to `(batch, out_dim)` by evaluating
//! learned B-spline activations on each input dimension and combining them with
//! learnable coefficients plus a linear residual connection.
//!
//! Uses cuBLAS sgemm via `GpuTensor` for the two matmul operations.

use std::sync::Arc;

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

use crate::gpu_tensor::{gpu_add, gpu_matmul, GpuTensor};

use super::spline::BSplineBasis;

/// A single Kolmogorov-Arnold Network layer.
///
/// Forward pass:
/// 1. Expand each input dim through B-spline basis: `(batch, in_dim)` -> `(batch, in_dim * num_bases)`
/// 2. Multiply by coefficients: `(batch, in_dim * num_bases) @ coefficients(in_dim * num_bases, out_dim)` -> `(batch, out_dim)`
/// 3. Add residual: `input @ residual_weight(in_dim, out_dim)` -> `(batch, out_dim)`
#[derive(Debug)]
pub struct KANLayer {
    /// B-spline basis evaluator (shared across all input dims)
    spline_basis: BSplineBasis,
    /// Learnable spline coefficients: (`in_dim` * `num_bases`, `out_dim`)
    pub coefficients: GpuTensor,
    /// Learnable residual weights: (`in_dim`, `out_dim`)
    pub residual_weight: GpuTensor,
    /// Input dimension
    in_dim: usize,
    /// Number of basis functions
    num_bases: usize,
}

impl KANLayer {
    /// Create a new KAN layer with Xavier-initialized weights.
    pub fn new(
        in_dim: usize,
        out_dim: usize,
        grid_size: usize,
        spline_order: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let mut spline_basis = BSplineBasis::new(grid_size, spline_order, -2.0, 2.0)?;
        let num_bases = spline_basis.num_bases();

        // Xavier-like init
        let scale = (2.0 / (in_dim * num_bases + out_dim) as f64).sqrt() as f32;
        let coefficients =
            GpuTensor::randn(&[in_dim * num_bases, out_dim], scale, stream)?;

        let res_scale = (1.0 / in_dim as f64).sqrt() as f32;
        let residual_weight = GpuTensor::randn(&[in_dim, out_dim], res_scale, stream)?;

        // Pre-compute B-spline lookup table
        spline_basis.precompute_gpu_grid(1024, stream)?;

        Ok(Self {
            spline_basis,
            coefficients,
            residual_weight,
            in_dim,
            num_bases,
        })
    }

    /// Forward pass: `(batch, in_dim)` -> `(batch, out_dim)`.
    pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor, MLError> {
        let batch = input.dim(0)?;

        // Flatten input to (batch * in_dim,) for basis evaluation
        let flat_input = input.reshape(&[batch * self.in_dim])?;

        // Evaluate B-spline basis: (batch * in_dim, num_bases)
        let basis_vals = self.spline_basis.evaluate(&flat_input)?;

        // Reshape to (batch, in_dim * num_bases)
        let basis_flat = basis_vals.reshape(&[batch, self.in_dim * self.num_bases])?;

        // Spline contribution: (batch, in_dim * num_bases) @ (in_dim * num_bases, out_dim)
        let spline_out = gpu_matmul(&basis_flat, &self.coefficients)?;

        // Residual contribution: (batch, in_dim) @ (in_dim, out_dim)
        let residual_out = gpu_matmul(input, &self.residual_weight)?;

        // Combine
        gpu_add(&spline_out, &residual_out)
    }
}

#[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")
    }

    #[test]
    fn test_kan_layer_output_shape() {
        let stream = test_stream();
        let layer = KANLayer::new(4, 8, 5, 4, &stream).unwrap();

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

    #[test]
    fn test_kan_layer_single_sample() {
        let stream = test_stream();
        let layer = KANLayer::new(4, 2, 3, 3, &stream).unwrap();
        let input = GpuTensor::from_vec(vec![1.0, 0.5, -0.5, 0.0], &[1, 4], &stream).unwrap();
        let output = layer.forward(&input).unwrap();
        assert_eq!(output.shape, vec![1, 2]);
        let v = output.to_vec().unwrap(); // test-only readback
        for val in &v {
            assert!(val.is_finite(), "Non-finite output: {}", val);
        }
    }
}