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.

use candle_core::Tensor;
use candle_nn::VarBuilder;
use ml_core::MLError;

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)
    coefficients: Tensor,
    /// Learnable residual weights: (in_dim, out_dim)
    residual_weight: Tensor,
    /// Input dimension
    in_dim: usize,
    /// Number of basis functions
    num_bases: usize,
}

impl KANLayer {
    /// Create a new KAN layer.
    ///
    /// # Arguments
    /// * `in_dim` - Input dimension
    /// * `out_dim` - Output dimension
    /// * `grid_size` - B-spline grid points
    /// * `spline_order` - B-spline order (degree + 1)
    /// * `vb` - VarBuilder for creating learnable parameters
    pub fn new(
        in_dim: usize,
        out_dim: usize,
        grid_size: usize,
        spline_order: usize,
        vb: VarBuilder<'_>,
    ) -> 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 scale
        let scale = (2.0 / (in_dim * num_bases + out_dim) as f64).sqrt();

        let coefficients = vb
            .get_with_hints(
                (in_dim * num_bases, out_dim),
                "coefficients",
                candle_nn::Init::Randn { mean: 0.0, stdev: scale },
            )
            .map_err(|e| MLError::ModelError(format!("Failed to create coefficients: {}", e)))?;

        let residual_weight = vb
            .get_with_hints(
                (in_dim, out_dim),
                "residual",
                candle_nn::Init::Randn {
                    mean: 0.0,
                    stdev: (1.0 / in_dim as f64).sqrt(),
                },
            )
            .map_err(|e| MLError::ModelError(format!("Failed to create residual weight: {}", e)))?;

        // Pre-compute B-spline lookup table so evaluate() uses GPU gather + lerp
        // instead of O(n * num_bases * 2^(k-1)) recursive CPU Cox-de Boor calls.
        // Infer device from the coefficients tensor we just created.
        let device = coefficients.device().clone();
        spline_basis.precompute_gpu_grid(1024, &device)?;

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

    /// Forward pass: `(batch, in_dim)` -> `(batch, out_dim)`.
    pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
        let device = input.device();
        let dims = input.shape().dims();
        let batch = *dims.first().ok_or_else(|| {
            MLError::ModelError("Input tensor has no dimensions".to_owned())
        })?;

        // Flatten input to (batch * in_dim,) for basis evaluation
        let flat_input = input.reshape(&[batch * self.in_dim]).map_err(|e| {
            MLError::ModelError(format!("KAN layer reshape failed: {}", e))
        })?;

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

        // Reshape to (batch, in_dim * num_bases)
        let basis_flat = basis_vals
            .reshape(&[batch, self.in_dim * self.num_bases])
            .map_err(|e| {
                MLError::ModelError(format!("KAN basis reshape failed: {}", e))
            })?;

        // Cast basis_flat to match coefficients dtype (BSpline evaluates in F32
        // for precision, but coefficients are in training dtype from VarBuilder)
        let coeff_dtype = self.coefficients.dtype();
        let basis_flat = basis_flat.to_dtype(coeff_dtype).map_err(|e| {
            MLError::ModelError(format!("KAN basis dtype cast failed: {}", e))
        })?;

        // Spline contribution: (batch, in_dim * num_bases) @ (in_dim * num_bases, out_dim)
        let spline_out = basis_flat.matmul(&self.coefficients).map_err(|e| {
            MLError::ModelError(format!("KAN spline matmul failed: {}", e))
        })?;

        // Cast input to match residual_weight dtype for matmul compatibility
        let input_cast = input.to_dtype(coeff_dtype).map_err(|e| {
            MLError::ModelError(format!("KAN input dtype cast failed: {}", e))
        })?;

        // Residual contribution: (batch, in_dim) @ (in_dim, out_dim)
        let residual_out = input_cast.matmul(&self.residual_weight).map_err(|e| {
            MLError::ModelError(format!("KAN residual matmul failed: {}", e))
        })?;

        // Combine
        spline_out.add(&residual_out).map_err(|e| {
            MLError::ModelError(format!("KAN layer add failed: {}", e))
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::Device;
    use candle_nn::VarMap;
    use ml_core::mixed_precision::training_dtype;

    #[test]
    fn test_kan_layer_output_shape() {
        let var_map = VarMap::new();
        let vb = VarBuilder::from_varmap(&var_map, training_dtype(&Device::Cpu), &Device::Cpu);
        let layer = KANLayer::new(4, 8, 5, 4, vb.pp("layer0")).unwrap();

        let input = Tensor::randn(0.0_f32, 0.5, &[2, 4], &Device::Cpu).unwrap();
        let output = layer.forward(&input).unwrap();
        let dims = output.shape().dims();
        assert_eq!(dims, &[2, 8]);
    }

    #[test]
    fn test_kan_layer_learnable_params() {
        let var_map = VarMap::new();
        let vb = VarBuilder::from_varmap(&var_map, training_dtype(&Device::Cpu), &Device::Cpu);
        let _layer = KANLayer::new(4, 8, 5, 4, vb.pp("layer0")).unwrap();

        let all_vars = var_map.all_vars();
        // Should have coefficients + residual = 2 tensors
        assert_eq!(all_vars.len(), 2, "Expected 2 learnable tensors, got {}", all_vars.len());
    }
}