gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ml): add KAN architecture + TLOB/KAN trainable/hyperopt adapters

Browse Source 
Phase 2-3 of ensemble expansion:
- KAN (Kolmogorov-Arnold Network): B-spline basis, layer, network, trainable adapter
- TLOB UnifiedTrainable adapter with 3D input support (batch, seq, features)
- Hyperopt adapters for both KAN and TLOB (ParameterSpace + metrics)
- ModelType::KAN variant registered in common, coordinator, lib.rs
- 44 new tests, all passing, zero warnings

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-23 01:32:25 +01:00
parent d777d714b2
commit 4eacd4e22f
14 changed files with 2146 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
common/src/model_types.rs
View File

@@ -42,6 +42,8 @@ pub enum ModelType {
TGNN,
/// Ensemble of multiple models
Ensemble,
/// Kolmogorov-Arnold Network
KAN,
}
impl ModelType {
@@ -60,6 +62,7 @@ impl ModelType {
Self::PPO => "ppo",
Self::Transformer => "transformer",
Self::Ensemble => "ensemble",
Self::KAN => "kan",
}
}
@@ -76,6 +79,7 @@ impl ModelType {
Self::DistilledMicroNet => "distilled",
Self::Transformer => "transformer",
Self::Ensemble => "ensemble",
Self::KAN => "kan",
}
}
@@ -92,6 +96,7 @@ impl ModelType {
Self::TGGN | Self::TGNN => "tggn",
Self::DistilledMicroNet => "distilled",
Self::Transformer => "transformer",
Self::KAN => "kan",
}
}
@@ -110,6 +115,7 @@ impl ModelType {
Self::TLOB => "TLOB",
Self::Transformer => "TRANSFORMER",
Self::Ensemble => "ENSEMBLE",
Self::KAN => "KAN",
}
}
@@ -134,6 +140,7 @@ impl ModelType {
"ppo" => Some(Self::PPO),
"transformer" => Some(Self::Transformer),
"ensemble" => Some(Self::Ensemble),
"kan" => Some(Self::KAN),
_ => None,
}
}

217
ml/src/hyperopt/adapters/kan.rs Normal file
View File

@@ -0,0 +1,217 @@
//! KAN (Kolmogorov-Arnold Network) Hyperparameter Optimization Adapter
//!
//! Provides `KANParams` (implementing `ParameterSpace`) and `KANMetrics`
//! for hyperparameter optimization of the KAN model via the unified
//! optimization framework.
use crate::hyperopt::traits::ParameterSpace;
use crate::MLError;
use serde::{Deserialize, Serialize};
/// Hyperparameters for KAN hyperopt tuning.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct KANParams {
/// Learning rate for AdamW optimizer (log-scale)
pub learning_rate: f64,
/// Number of B-spline grid points per activation
pub grid_size: usize,
/// B-spline order (degree + 1)
pub spline_order: usize,
/// Hidden layer width (used for all hidden layers)
pub hidden_width: usize,
/// Number of hidden layers (total layers = num_layers + 1 for output)
pub num_layers: usize,
/// L2 weight decay (log-scale)
pub weight_decay: f64,
/// Maximum gradient norm for clipping (log-scale)
pub grad_clip: f64,
/// Training batch size
pub batch_size: usize,
}
impl Default for KANParams {
fn default() -> Self {
Self {
learning_rate: 1e-3,
grid_size: 5,
spline_order: 4,
hidden_width: 32,
num_layers: 2,
weight_decay: 1e-4,
grad_clip: 1.0,
batch_size: 32,
}
}
}
/// Number of continuous parameters in the KAN parameter space.
const NUM_PARAMS: usize = 8;
impl ParameterSpace for KANParams {
fn continuous_bounds() -> Vec<(f64, f64)> {
vec![
(1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log)
(3.0, 12.0), // grid_size (linear)
(2.0, 5.0), // spline_order (linear)
(16.0, 64.0), // hidden_width (linear)
(2.0, 4.0), // num_layers (linear)
(1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log)
(0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log)
(8.0, 64.0), // batch_size (linear)
]
}
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
if x.len() != NUM_PARAMS {
return Err(MLError::ConfigError {
reason: format!("Expected {} continuous parameters, got {}", NUM_PARAMS, x.len()),
});
}
Ok(Self {
learning_rate: x.first().copied().unwrap_or(0.0).exp(),
grid_size: x.get(1).copied().unwrap_or(5.0).round().max(3.0) as usize,
spline_order: x.get(2).copied().unwrap_or(4.0).round().max(2.0) as usize,
hidden_width: x.get(3).copied().unwrap_or(32.0).round().max(16.0) as usize,
num_layers: x.get(4).copied().unwrap_or(2.0).round().max(2.0) as usize,
weight_decay: x.get(5).copied().unwrap_or(0.0).exp(),
grad_clip: x.get(6).copied().unwrap_or(0.0).exp(),
batch_size: x.get(7).copied().unwrap_or(32.0).round().max(8.0) as usize,
})
}
fn to_continuous(&self) -> Vec<f64> {
vec![
self.learning_rate.ln(),
self.grid_size as f64,
self.spline_order as f64,
self.hidden_width as f64,
self.num_layers as f64,
self.weight_decay.ln(),
self.grad_clip.ln(),
self.batch_size as f64,
]
}
fn param_names() -> Vec<&'static str> {
vec![
"learning_rate",
"grid_size",
"spline_order",
"hidden_width",
"num_layers",
"weight_decay",
"grad_clip",
"batch_size",
]
}
}
/// Metrics returned from KAN hyperopt training.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct KANMetrics {
/// Validation loss (primary objective for optimization)
pub val_loss: f64,
/// Training loss at end of run
pub train_loss: f64,
/// Number of epochs completed
pub epochs_completed: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bounds_count_matches_param_names() {
let bounds = KANParams::continuous_bounds();
let names = KANParams::param_names();
assert_eq!(bounds.len(), names.len());
assert_eq!(bounds.len(), NUM_PARAMS);
}
#[test]
fn test_roundtrip_continuous() {
let params = KANParams::default();
let continuous = params.to_continuous();
let restored = KANParams::from_continuous(&continuous).unwrap();
assert!(
(params.learning_rate - restored.learning_rate).abs() < 1e-6,
"lr mismatch: {} vs {}",
params.learning_rate,
restored.learning_rate
);
assert_eq!(params.grid_size, restored.grid_size);
assert_eq!(params.spline_order, restored.spline_order);
assert_eq!(params.hidden_width, restored.hidden_width);
assert_eq!(params.num_layers, restored.num_layers);
assert!(
(params.weight_decay - restored.weight_decay).abs() / params.weight_decay < 1e-6,
"weight_decay mismatch"
);
assert!(
(params.grad_clip - restored.grad_clip).abs() < 1e-6,
"grad_clip mismatch"
);
assert_eq!(params.batch_size, restored.batch_size);
}
#[test]
fn test_from_continuous_wrong_length_errors() {
let result = KANParams::from_continuous(&[0.1, 0.2]);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(
err_msg.contains("Expected 8"),
"Error should mention expected count: {}",
err_msg
);
}
#[test]
fn test_bounds_are_valid() {
for (i, (min, max)) in KANParams::continuous_bounds().iter().enumerate() {
assert!(
min < max,
"Invalid bounds for param {} ({}): {} >= {}",
i,
KANParams::param_names().get(i).copied().unwrap_or("?"),
min,
max
);
}
}
#[test]
fn test_default_within_bounds() {
let params = KANParams::default();
let continuous = params.to_continuous();
let bounds = KANParams::continuous_bounds();
let names = KANParams::param_names();
for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() {
assert!(
*val >= *min && *val <= *max,
"Param {} ({}) = {} outside [{}, {}]",
i,
names.get(i).copied().unwrap_or("?"),
val,
min,
max
);
}
}
#[test]
fn test_metrics_default() {
let metrics = KANMetrics::default();
assert_eq!(metrics.val_loss, 0.0);
assert_eq!(metrics.train_loss, 0.0);
assert_eq!(metrics.epochs_completed, 0);
}
#[test]
fn test_continuous_length() {
let params = KANParams::default();
let continuous = params.to_continuous();
assert_eq!(continuous.len(), NUM_PARAMS);
}
}

4
ml/src/hyperopt/adapters/mod.rs
View File

@@ -52,14 +52,17 @@
pub mod async_data_loader;
pub mod continuous_ppo;
pub mod dqn;
pub mod kan;
pub mod liquid;
pub mod mamba2;
pub mod ppo;
pub mod tft;
pub mod tggn;
pub mod tlob;
// Re-export adapters for convenience
pub use async_data_loader::AsyncDataLoader;
pub use kan::{KANMetrics, KANParams};
pub use liquid::LiquidParams;
pub use continuous_ppo::{ContinuousPPOMetrics, ContinuousPPOParams, ContinuousPPOTrainer};
pub use dqn::{DQNMetrics, DQNParams, DQNTrainer};
@@ -67,3 +70,4 @@ pub use mamba2::{Mamba2Metrics, Mamba2Params, Mamba2Trainer};
pub use ppo::{PPOMetrics, PPOParams, PPOTrainer};
pub use tft::{TFTMetrics, TFTParams, TFTTrainer as TFTHyperoptTrainer};
pub use tggn::{TGGNMetrics, TGGNParams};
pub use tlob::{TLOBMetrics, TLOBParams};

229
ml/src/hyperopt/adapters/tlob.rs Normal file
View File

@@ -0,0 +1,229 @@
//! TLOB (Time Limit Order Book) Hyperparameter Optimization Adapter
//!
//! Provides `TLOBParams` (implementing `ParameterSpace`) and `TLOBMetrics`
//! for hyperparameter optimization of the TLOB model via the unified
//! optimization framework.
use crate::hyperopt::traits::ParameterSpace;
use crate::MLError;
use serde::{Deserialize, Serialize};
/// Hyperparameters for TLOB hyperopt tuning.
///
/// Controls the TLOB transformer's architecture and training:
/// - Architecture: `d_model`, `num_heads`, `num_layers`, `seq_len`
/// - Regularization: `dropout`, `weight_decay`, `grad_clip`
/// - Training: `learning_rate`, `batch_size`
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TLOBParams {
/// Learning rate for AdamW optimizer (log-scale)
pub learning_rate: f64,
/// Model / hidden dimension for projection layers
pub d_model: usize,
/// Number of attention heads
pub num_heads: usize,
/// Number of transformer/projection layers
pub num_layers: usize,
/// Sequence length (number of time steps in LOB snapshot)
pub seq_len: usize,
/// Dropout rate for regularization
pub dropout: f64,
/// Training batch size
pub batch_size: usize,
/// L2 weight decay (log-scale)
pub weight_decay: f64,
/// Maximum gradient norm for clipping (log-scale)
pub grad_clip: f64,
}
impl Default for TLOBParams {
fn default() -> Self {
Self {
learning_rate: 1e-3,
d_model: 128,
num_heads: 4,
num_layers: 2,
seq_len: 128,
dropout: 0.1,
batch_size: 32,
weight_decay: 1e-4,
grad_clip: 1.0,
}
}
}
/// Number of continuous parameters in the TLOB parameter space.
const NUM_PARAMS: usize = 9;
impl ParameterSpace for TLOBParams {
fn continuous_bounds() -> Vec<(f64, f64)> {
vec![
(1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log)
(64.0, 512.0), // d_model (linear)
(2.0, 16.0), // num_heads (linear)
(1.0, 8.0), // num_layers (linear)
(32.0, 256.0), // seq_len (linear)
(0.0, 0.5), // dropout (linear)
(8.0, 64.0), // batch_size (linear)
(1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log)
(0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log)
]
}
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
if x.len() != NUM_PARAMS {
return Err(MLError::ConfigError {
reason: format!("Expected {} continuous parameters, got {}", NUM_PARAMS, x.len()),
});
}
Ok(Self {
learning_rate: x.first().copied().unwrap_or(0.0).exp(),
d_model: x.get(1).copied().unwrap_or(128.0).round().max(64.0) as usize,
num_heads: x.get(2).copied().unwrap_or(4.0).round().max(2.0) as usize,
num_layers: x.get(3).copied().unwrap_or(2.0).round().max(1.0) as usize,
seq_len: x.get(4).copied().unwrap_or(128.0).round().max(32.0) as usize,
dropout: x.get(5).copied().unwrap_or(0.1).clamp(0.0, 0.5),
batch_size: x.get(6).copied().unwrap_or(32.0).round().max(8.0) as usize,
weight_decay: x.get(7).copied().unwrap_or(0.0).exp(),
grad_clip: x.get(8).copied().unwrap_or(0.0).exp(),
})
}
fn to_continuous(&self) -> Vec<f64> {
vec![
self.learning_rate.ln(),
self.d_model as f64,
self.num_heads as f64,
self.num_layers as f64,
self.seq_len as f64,
self.dropout,
self.batch_size as f64,
self.weight_decay.ln(),
self.grad_clip.ln(),
]
}
fn param_names() -> Vec<&'static str> {
vec![
"learning_rate",
"d_model",
"num_heads",
"num_layers",
"seq_len",
"dropout",
"batch_size",
"weight_decay",
"grad_clip",
]
}
}
/// Metrics returned from TLOB hyperopt training.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TLOBMetrics {
/// Validation loss (primary objective for optimization)
pub val_loss: f64,
/// Training loss at end of run
pub train_loss: f64,
/// Directional accuracy on validation set
pub directional_accuracy: f64,
/// Number of epochs completed
pub epochs_completed: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bounds_count_matches_param_names() {
let bounds = TLOBParams::continuous_bounds();
let names = TLOBParams::param_names();
assert_eq!(bounds.len(), names.len());
assert_eq!(bounds.len(), NUM_PARAMS);
}
#[test]
fn test_roundtrip_continuous() {
let params = TLOBParams::default();
let continuous = params.to_continuous();
let restored = TLOBParams::from_continuous(&continuous).unwrap();
assert!(
(params.learning_rate - restored.learning_rate).abs() < 1e-6,
"lr mismatch: {} vs {}",
params.learning_rate,
restored.learning_rate
);
assert_eq!(params.d_model, restored.d_model);
assert_eq!(params.num_heads, restored.num_heads);
assert_eq!(params.num_layers, restored.num_layers);
assert_eq!(params.seq_len, restored.seq_len);
assert!((params.dropout - restored.dropout).abs() < 1e-6, "dropout mismatch");
assert_eq!(params.batch_size, restored.batch_size);
assert!(
(params.weight_decay - restored.weight_decay).abs() / params.weight_decay < 1e-6,
"weight_decay mismatch"
);
assert!(
(params.grad_clip - restored.grad_clip).abs() < 1e-6,
"grad_clip mismatch"
);
}
#[test]
fn test_from_continuous_wrong_length_errors() {
let result = TLOBParams::from_continuous(&[0.1, 0.2]);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(err_msg.contains("Expected 9"), "Error should mention expected count: {}", err_msg);
}
#[test]
fn test_bounds_are_valid() {
for (i, (min, max)) in TLOBParams::continuous_bounds().iter().enumerate() {
assert!(
min < max,
"Invalid bounds for param {} ({}): {} >= {}",
i,
TLOBParams::param_names().get(i).copied().unwrap_or("?"),
min,
max
);
}
}
#[test]
fn test_default_within_bounds() {
let params = TLOBParams::default();
let continuous = params.to_continuous();
let bounds = TLOBParams::continuous_bounds();
let names = TLOBParams::param_names();
for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() {
assert!(
*val >= *min && *val <= *max,
"Param {} ({}) = {} outside [{}, {}]",
i,
names.get(i).copied().unwrap_or("?"),
val,
min,
max
);
}
}
#[test]
fn test_metrics_default() {
let metrics = TLOBMetrics::default();
assert_eq!(metrics.val_loss, 0.0);
assert_eq!(metrics.train_loss, 0.0);
assert_eq!(metrics.directional_accuracy, 0.0);
assert_eq!(metrics.epochs_completed, 0);
}
#[test]
fn test_continuous_length() {
let params = TLOBParams::default();
let continuous = params.to_continuous();
assert_eq!(continuous.len(), NUM_PARAMS);
}
}

5
ml/src/integration/coordinator.rs
View File

@@ -471,6 +471,10 @@ impl EnsembleCoordinator {
// Ensemble methods - use weighted combination approach
self.deep_q_prediction(input_features, model.weight * 1.2) // Enhanced prediction for ensemble
},
ModelType::KAN => {
// Kolmogorov-Arnold Network prediction
self.deep_q_prediction(input_features, model.weight)
},
};
let latency_us = start_time.elapsed().as_micros() as u64;
@@ -882,6 +886,7 @@ impl EnsembleCoordinator {
ModelType::PPO => 0.83, // Good confidence for policy optimization
ModelType::Transformer => 0.88, // High confidence for sequence modeling
ModelType::Ensemble => 0.92, // Highest confidence for ensemble methods
ModelType::KAN => 0.86, // Good confidence for learned activation functions
};
// Adjust confidence based on feature quality

42
ml/src/kan/config.rs Normal file
View File

@@ -0,0 +1,42 @@
//! KAN (Kolmogorov-Arnold Network) configuration.
//!
//! Defines hyperparameters for KAN models including B-spline grid size,
//! spline order, layer widths, and training parameters.
use serde::{Deserialize, Serialize};
/// Configuration for a Kolmogorov-Arnold Network.
///
/// KAN replaces fixed activation functions (ReLU) with learnable B-spline
/// activations on each edge, enabling the network to discover arbitrary
/// non-linear relationships.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KANConfig {
/// Number of B-spline grid points per activation function.
/// More points = finer approximation but more parameters.
pub grid_size: usize,
/// B-spline order (degree + 1). 4 = cubic splines.
pub spline_order: usize,
/// Layer widths including input and output dimensions.
/// Example: `[51, 32, 16, 1]` means input=51, two hidden layers, output=1.
pub layer_widths: Vec<usize>,
/// Learning rate for AdamW optimizer.
pub learning_rate: f64,
/// L2 weight decay for regularization.
pub weight_decay: f64,
/// Maximum gradient norm for clipping.
pub grad_clip: f64,
}
impl Default for KANConfig {
fn default() -> Self {
Self {
grid_size: 5,
spline_order: 4,
layer_widths: vec![51, 32, 16, 1],
learning_rate: 1e-3,
weight_decay: 1e-4,
grad_clip: 1.0,
}
}
}

151
ml/src/kan/layer.rs Normal file
View File

@@ -0,0 +1,151 @@
//! 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 crate::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 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)))?;
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_string())
})?;
// 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))
})?;
// 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))
})?;
// Residual contribution: (batch, in_dim) @ (in_dim, out_dim)
let residual_out = input.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::{DType, Device};
use candle_nn::VarMap;
#[test]
fn test_kan_layer_output_shape() {
let var_map = VarMap::new();
let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu);
let layer = KANLayer::new(4, 8, 5, 4, vb.pp("layer0")).unwrap();
let input = Tensor::randn(0.0f32, 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, DType::F32, &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());
}
}

15
ml/src/kan/mod.rs Normal file
View File

@@ -0,0 +1,15 @@
//! KAN (Kolmogorov-Arnold Network) module.
//!
//! Implements KAN with learnable B-spline activation functions on each edge
//! of the network, replacing fixed activations (ReLU) with data-driven
//! non-linearities.
pub mod config;
pub mod layer;
pub mod network;
pub mod spline;
pub mod trainable;
pub use config::KANConfig;
pub use network::KANNetwork;
pub use trainable::KANTrainableAdapter;

121
ml/src/kan/network.rs Normal file
View File

@@ -0,0 +1,121 @@
//! 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 candle_core::Tensor;
use candle_nn::VarBuilder;
use crate::MLError;
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`.
/// For example, widths `[51, 32, 16, 1]` creates 3 layers:
/// 51->32, 32->16, 16->1.
pub fn new(config: &KANConfig, vb: VarBuilder<'_>) -> Result<Self, MLError> {
if config.layer_widths.len() < 2 {
return Err(MLError::ConfigError {
reason: "layer_widths must have at least 2 entries".to_string(),
});
}
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 {
reason: format!("Missing layer_widths index {}", i),
})?;
let out_dim = *config.layer_widths.get(i + 1).ok_or_else(|| MLError::ConfigError {
reason: format!("Missing layer_widths index {}", i + 1),
})?;
let layer = KANLayer::new(
in_dim,
out_dim,
config.grid_size,
config.spline_order,
vb.pp(format!("layer_{}", i)),
)?;
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: &Tensor) -> Result<Tensor, MLError> {
let mut x = input.clone();
for layer in &self.layers {
x = layer.forward(&x)?;
}
Ok(x)
}
}
#[cfg(test)]
mod tests {
use super::*;
use candle_core::{DType, Device};
use candle_nn::VarMap;
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 config = make_config();
let var_map = VarMap::new();
let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu);
let net = KANNetwork::new(&config, vb).unwrap();
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let output = net.forward(&input).unwrap();
let dims = output.shape().dims();
assert_eq!(dims, &[2, 1], "Expected (2, 1), got {:?}", dims);
}
#[test]
fn test_kan_network_produces_gradients() {
let config = make_config();
let var_map = VarMap::new();
let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu);
let net = KANNetwork::new(&config, vb).unwrap();
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let output = net.forward(&input).unwrap();
let loss = output.sqr().unwrap().mean_all().unwrap();
let grads = loss.backward().unwrap();
// Check that at least some variables received gradients
let vars_lock = var_map.data().lock().unwrap();
let mut has_grad = false;
for (_name, var) in vars_lock.iter() {
if grads.get(var.as_tensor()).is_some() {
has_grad = true;
break;
}
}
assert!(has_grad, "No gradients found for any variable");
}
}

188
ml/src/kan/spline.rs Normal file
View File

@@ -0,0 +1,188 @@
//! B-spline basis function evaluator for KAN.
//!
//! Implements Cox-de Boor recursion on CPU to compute B-spline basis values.
//! Knots are stored as `Vec<f32>` (no gradient needed for the knot vector).
//! The evaluator produces a `(batch, num_bases)` tensor for use in KAN layers.
use candle_core::{DType, Device, Tensor};
use crate::MLError;
/// B-spline basis evaluator.
///
/// Stores a uniform knot vector and evaluates all basis functions for a
/// batch of scalar inputs using the Cox-de Boor recursion.
#[derive(Debug, Clone)]
pub struct BSplineBasis {
/// Uniform knot vector (length = grid_size + 2 * order).
knots: Vec<f32>,
/// Spline order (degree + 1).
order: usize,
/// Number of basis functions = grid_size + order - 1.
num_bases: usize,
}
impl BSplineBasis {
/// Create a new B-spline basis with uniform knots over `[x_min, x_max]`.
///
/// # Arguments
/// * `grid_size` - Number of interior grid points
/// * `order` - Spline order (degree + 1), e.g. 4 for cubic
/// * `x_min` - Left boundary of the domain
/// * `x_max` - Right boundary of the domain
pub fn new(grid_size: usize, order: usize, x_min: f64, x_max: f64) -> Result<Self, MLError> {
if grid_size == 0 || order == 0 {
return Err(MLError::ConfigError {
reason: "grid_size and order must be > 0".to_string(),
});
}
if x_max <= x_min {
return Err(MLError::ConfigError {
reason: format!("x_max ({}) must be > x_min ({})", x_max, x_min),
});
}
let num_bases = grid_size + order - 1;
// Total knots: num_bases + order = grid_size + 2*order - 1
let n_knots = num_bases + order;
let mut knots = Vec::with_capacity(n_knots);
// Extend domain slightly beyond [x_min, x_max] for boundary padding.
let step = (x_max - x_min) / grid_size as f64;
let pad = (order as f64 - 1.0) * step;
let start = x_min - pad;
for i in 0..n_knots {
knots.push((start + i as f64 * step) as f32);
}
Ok(Self {
knots,
order,
num_bases,
})
}
/// Number of basis functions produced.
pub fn num_bases(&self) -> usize {
self.num_bases
}
/// Evaluate all basis functions for a batch of scalars.
///
/// # Arguments
/// * `x` - Tensor of shape `(n,)` containing scalar inputs
/// * `device` - Device to place the output on
///
/// # Returns
/// Tensor of shape `(n, num_bases)` with basis values.
pub fn evaluate(&self, x: &Tensor, device: &Device) -> Result<Tensor, MLError> {
// Flatten to 1-D regardless of input shape
let x_flat = x.flatten_all().map_err(|e| {
MLError::ModelError(format!("BSpline flatten failed: {}", e))
})?;
let n = x_flat.dims1().map_err(|e| {
MLError::ModelError(format!("BSpline dims failed: {}", e))
})?;
let x_vec: Vec<f32> = x_flat.to_vec1().map_err(|e| {
MLError::ModelError(format!("BSpline to_vec1 failed: {}", e))
})?;
// Compute basis values on CPU via Cox-de Boor
let mut result = vec![0.0f32; n * self.num_bases];
for (i, &xi) in x_vec.iter().enumerate() {
for j in 0..self.num_bases {
let val = self.cox_de_boor(xi, j, self.order);
// Safe index: i * num_bases + j is always within bounds
if let Some(slot) = result.get_mut(i * self.num_bases + j) {
*slot = val;
}
}
}
Tensor::from_vec(result, &[n, self.num_bases], device)
.map_err(|e| MLError::ModelError(format!("BSpline tensor creation failed: {}", e)))?
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("BSpline dtype conversion failed: {}", e)))
}
/// Cox-de Boor recursion for basis function `j` of given `order` at point `x`.
fn cox_de_boor(&self, x: f32, j: usize, k: usize) -> f32 {
if k == 1 {
let t_j = self.knots.get(j).copied().unwrap_or(0.0);
let t_j1 = self.knots.get(j + 1).copied().unwrap_or(0.0);
return if x >= t_j && x < t_j1 { 1.0 } else { 0.0 };
}
let t_j = self.knots.get(j).copied().unwrap_or(0.0);
let t_jk = self.knots.get(j + k).copied().unwrap_or(0.0);
let t_jk1 = self.knots.get(j + k - 1).copied().unwrap_or(0.0);
let t_j1 = self.knots.get(j + 1).copied().unwrap_or(0.0);
let denom_left = t_jk1 - t_j;
let left = if denom_left.abs() > 1e-10 {
((x - t_j) / denom_left) * self.cox_de_boor(x, j, k - 1)
} else {
0.0
};
let denom_right = t_jk - t_j1;
let right = if denom_right.abs() > 1e-10 {
((t_jk - x) / denom_right) * self.cox_de_boor(x, j + 1, k - 1)
} else {
0.0
};
left + right
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bspline_basis_shape() {
let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
assert_eq!(basis.num_bases(), 8); // 5 + 4 - 1
let x = Tensor::new(&[0.0f32, 0.5, -0.5, 0.1], &Device::Cpu).unwrap();
let out = basis.evaluate(&x, &Device::Cpu).unwrap();
let dims = out.shape().dims();
assert_eq!(dims.len(), 2);
assert_eq!(dims[0], 4);
assert_eq!(dims[1], 8);
}
#[test]
fn test_bspline_partition_of_unity() {
let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
// Points well inside the domain should sum close to 1.0
let x = Tensor::new(&[0.0f32, 0.3, -0.3], &Device::Cpu).unwrap();
let out = basis.evaluate(&x, &Device::Cpu).unwrap();
let sums = out.sum(1).unwrap();
let sums_vec: Vec<f32> = sums.to_vec1().unwrap();
for (i, &s) in sums_vec.iter().enumerate() {
assert!(
(s - 1.0).abs() < 0.15,
"Basis sum at index {} = {} (expected ~1.0)",
i,
s
);
}
}
#[test]
fn test_bspline_non_negative() {
let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
let x = Tensor::new(&[0.0f32, 0.5, -0.5, 0.9], &Device::Cpu).unwrap();
let out = basis.evaluate(&x, &Device::Cpu).unwrap();
let flat: Vec<f32> = out.flatten_all().unwrap().to_vec1().unwrap();
for &v in &flat {
assert!(
v >= -1e-6,
"Basis value {} is negative (< -epsilon)",
v
);
}
}
}

493
ml/src/kan/trainable.rs Normal file
View File

@@ -0,0 +1,493 @@
//! UnifiedTrainable adapter for KAN (Kolmogorov-Arnold Network).
//!
//! Wraps a KANNetwork with VarMap + AdamW optimizer to provide the
//! UnifiedTrainable interface. Follows the same pattern as TGGNTrainableAdapter.
use candle_core::{backprop::GradStore, DType, Device, Tensor};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use std::collections::HashMap;
use super::config::KANConfig;
use super::network::KANNetwork;
use crate::training::unified_trainer::{checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable};
use crate::MLError;
/// KAN trainable adapter implementing UnifiedTrainable.
pub struct KANTrainableAdapter {
config: KANConfig,
var_map: VarMap,
network: KANNetwork,
optimizer: AdamW,
grads: Option<GradStore>,
device: Device,
learning_rate: f64,
step: usize,
latest_metrics: TrainingMetrics,
loss_history: Vec<f64>,
last_grad_norm: f64,
}
impl std::fmt::Debug for KANTrainableAdapter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("KANTrainableAdapter")
.field("config", &self.config)
.field("device", &format!("{:?}", self.device))
.field("learning_rate", &self.learning_rate)
.field("step", &self.step)
.field("loss_history_len", &self.loss_history.len())
.field("last_grad_norm", &self.last_grad_norm)
.finish_non_exhaustive()
}
}
impl KANTrainableAdapter {
/// Create a new KAN trainable adapter.
pub fn new(config: KANConfig, device: &Device) -> Result<Self, MLError> {
let var_map = VarMap::new();
let vb = VarBuilder::from_varmap(&var_map, DType::F32, device);
let network = KANNetwork::new(&config, vb)?;
let learning_rate = config.learning_rate;
let weight_decay = config.weight_decay;
let all_vars = var_map.all_vars();
let optimizer = AdamW::new(
all_vars,
ParamsAdamW {
lr: learning_rate,
beta1: 0.9,
beta2: 0.999,
eps: 1e-8,
weight_decay,
},
)
.map_err(|e| MLError::ModelError(format!("Failed to create AdamW optimizer: {}", e)))?;
Ok(Self {
config,
var_map,
network,
optimizer,
grads: None,
device: device.clone(),
learning_rate,
step: 0,
latest_metrics: TrainingMetrics::default(),
loss_history: Vec::new(),
last_grad_norm: 0.0,
})
}
/// Access the KAN configuration.
pub fn kan_config(&self) -> &KANConfig {
&self.config
}
}
impl UnifiedTrainable for KANTrainableAdapter {
fn model_type(&self) -> &str {
"KAN"
}
fn device(&self) -> &Device {
&self.device
}
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
self.network.forward(input)
}
fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor, MLError> {
let diff = predictions.sub(targets).map_err(|e| {
MLError::ModelError(format!("Loss subtraction failed: {}", e))
})?;
let squared = diff.powf(2.0).map_err(|e| {
MLError::ModelError(format!("Loss squaring failed: {}", e))
})?;
squared.mean_all().map_err(|e| {
MLError::ModelError(format!("Loss mean failed: {}", e))
})
}
fn backward(&mut self, loss: &Tensor) -> Result<f64, MLError> {
let grads = loss.backward().map_err(|e| {
MLError::TrainingError(format!("Backward pass failed: {}", e))
})?;
let mut grad_norm_sq = 0.0;
let vars_lock = self
.var_map
.data()
.lock()
.map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?;
for (_name, var) in vars_lock.iter() {
if let Some(grad) = grads.get(var.as_tensor()) {
let norm = grad
.sqr()
.and_then(|s| s.sum_all())
.and_then(|s| s.to_scalar::<f32>())
.map_err(|e| {
MLError::ModelError(format!("Failed to compute grad norm: {}", e))
})?;
grad_norm_sq += norm as f64;
}
}
drop(vars_lock);
let grad_norm = grad_norm_sq.sqrt();
self.last_grad_norm = grad_norm;
self.latest_metrics.grad_norm = Some(grad_norm);
self.grads = Some(grads);
Ok(grad_norm)
}
fn optimizer_step(&mut self) -> Result<(), MLError> {
if let Some(grads) = self.grads.take() {
self.optimizer.step(&grads).map_err(|e| {
MLError::TrainingError(format!("Optimizer step failed: {}", e))
})?;
}
self.step += 1;
Ok(())
}
fn zero_grad(&mut self) -> Result<(), MLError> {
Ok(())
}
fn get_learning_rate(&self) -> f64 {
self.learning_rate
}
fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
self.learning_rate = lr;
let all_vars = self.var_map.all_vars();
self.optimizer = AdamW::new(
all_vars,
ParamsAdamW {
lr,
beta1: 0.9,
beta2: 0.999,
eps: 1e-8,
weight_decay: self.config.weight_decay,
},
)
.map_err(|e| {
MLError::ModelError(format!("Failed to recreate optimizer with new lr: {}", e))
})?;
Ok(())
}
fn get_step(&self) -> usize {
self.step
}
fn collect_metrics(&self) -> TrainingMetrics {
let mut metrics = self.latest_metrics.clone();
if !self.loss_history.is_empty() {
let recent: Vec<f64> = self.loss_history.iter().rev().take(100).copied().collect();
metrics.loss = recent.iter().sum::<f64>() / recent.len() as f64;
}
metrics.learning_rate = self.learning_rate;
metrics.custom_metrics.insert("training_steps".to_string(), self.step as f64);
metrics.custom_metrics.insert("loss_history_len".to_string(), self.loss_history.len() as f64);
metrics.custom_metrics.insert("grid_size".to_string(), self.config.grid_size as f64);
metrics.custom_metrics.insert("spline_order".to_string(), self.config.spline_order as f64);
metrics.custom_metrics.insert("num_layers".to_string(), (self.config.layer_widths.len().saturating_sub(1)) as f64);
metrics
}
fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
let metadata = CheckpointMetadata {
model_type: "KAN".to_string(),
version: "1.0.0".to_string(),
epoch: 0,
step: self.step,
timestamp: std::time::SystemTime::now(),
config: serde_json::to_value(&self.config).map_err(|e| {
MLError::SerializationError {
reason: format!("Failed to serialize config: {}", e),
}
})?,
metrics: self.collect_metrics(),
};
checkpoint::save_metadata(&metadata, checkpoint_path)?;
let safetensors_path = format!("{}.safetensors", checkpoint_path);
let vars_lock = self
.var_map
.data()
.lock()
.map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?;
let mut tensors: HashMap<String, Tensor> = HashMap::new();
for (name, var) in vars_lock.iter() {
tensors.insert(name.clone(), var.as_tensor().clone());
}
drop(vars_lock);
candle_core::safetensors::save(&tensors, &safetensors_path).map_err(|e| {
MLError::CheckpointError(format!("Failed to save safetensors: {}", e))
})?;
tracing::info!("Saved KAN checkpoint to {} (step {})", checkpoint_path, self.step);
Ok(checkpoint_path.to_string())
}
fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
let metadata = checkpoint::load_metadata(checkpoint_path)?;
if metadata.model_type != "KAN" {
return Err(MLError::CheckpointError(format!(
"Invalid model type in checkpoint: expected KAN, got {}",
metadata.model_type
)));
}
let safetensors_path = format!("{}.safetensors", checkpoint_path);
let tensors = candle_core::safetensors::load(&safetensors_path, &self.device).map_err(
|e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)),
)?;
let vars_lock = self
.var_map
.data()
.lock()
.map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?;
for (name, tensor) in &tensors {
if let Some(var) = vars_lock.get(name) {
var.set(tensor).map_err(|e| {
MLError::CheckpointError(format!("Failed to set var {}: {}", name, e))
})?;
} else {
tracing::warn!("Checkpoint contains unknown variable: {}", name);
}
}
drop(vars_lock);
self.step = metadata.step;
self.latest_metrics = metadata.metrics.clone();
tracing::info!("Loaded KAN checkpoint from {} (step {})", checkpoint_path, metadata.step);
Ok(metadata)
}
fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result<f64, MLError> {
if val_data.is_empty() {
return Err(MLError::ValidationError {
message: "Empty validation dataset".to_string(),
});
}
let mut total_loss = 0.0;
let mut count = 0usize;
for (input, target) in val_data {
let prediction = self.forward(input)?;
let loss = self.compute_loss(&prediction, target)?;
let loss_val = loss.to_scalar::<f32>().map_err(|e| MLError::ValidationError {
message: format!("Failed to extract loss value: {}", e),
})?;
total_loss += loss_val as f64;
count += 1;
}
let avg_loss = if count > 0 {
total_loss / count as f64
} else {
0.0
};
self.latest_metrics.val_loss = Some(avg_loss);
tracing::debug!("KAN validation loss: {:.6}", avg_loss);
Ok(avg_loss)
}
}
#[cfg(test)]
mod tests {
use super::*;
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_model_type() {
let cfg = make_config();
let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
assert_eq!(adapter.model_type(), "KAN");
}
#[test]
fn test_device() {
let cfg = make_config();
let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
assert!(matches!(adapter.device(), &Device::Cpu));
}
#[test]
fn test_forward_shape() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let dims = output.shape().dims();
assert_eq!(dims.len(), 2);
assert_eq!(dims[0], 2);
assert_eq!(dims[1], 1);
}
#[test]
fn test_compute_loss() {
let cfg = make_config();
let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let preds = Tensor::new(&[[1.0f32], [2.0]], &Device::Cpu).unwrap();
let targets = Tensor::new(&[[1.5f32], [2.5]], &Device::Cpu).unwrap();
let loss = adapter.compute_loss(&preds, &targets).unwrap();
let loss_val: f32 = loss.to_scalar().unwrap();
assert!((loss_val - 0.25).abs() < 1e-5, "Expected ~0.25, got {}", loss_val);
}
#[test]
fn test_backward_returns_grad_norm() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let loss = adapter.compute_loss(&output, &targets).unwrap();
let grad_norm = adapter.backward(&loss).unwrap();
assert!(grad_norm >= 0.0, "Gradient norm should be non-negative");
assert!(adapter.grads.is_some(), "Grads should be stored after backward");
}
#[test]
fn test_train_step_cycle() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
assert_eq!(adapter.get_step(), 0);
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let loss = adapter.compute_loss(&output, &targets).unwrap();
let loss_val: f64 = loss.to_scalar::<f32>().unwrap() as f64;
adapter.backward(&loss).unwrap();
adapter.optimizer_step().unwrap();
assert_eq!(adapter.get_step(), 1);
adapter.loss_history.push(loss_val);
let metrics = adapter.collect_metrics();
assert!(metrics.loss >= 0.0);
assert_eq!(metrics.custom_metrics.get("training_steps").copied(), Some(1.0));
}
#[test]
fn test_learning_rate_get_set() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let original_lr = adapter.get_learning_rate();
assert!((original_lr - 1e-3).abs() < 1e-10);
adapter.set_learning_rate(5e-4).unwrap();
assert!((adapter.get_learning_rate() - 5e-4).abs() < 1e-10);
}
#[test]
fn test_collect_metrics() {
let cfg = make_config();
let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let metrics = adapter.collect_metrics();
assert!(metrics.custom_metrics.contains_key("training_steps"));
assert!(metrics.custom_metrics.contains_key("grid_size"));
assert!(metrics.custom_metrics.contains_key("spline_order"));
assert!(metrics.custom_metrics.contains_key("num_layers"));
assert_eq!(metrics.custom_metrics.get("grid_size").copied(), Some(3.0));
}
#[test]
fn test_checkpoint_roundtrip() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let loss = adapter.compute_loss(&output, &targets).unwrap();
adapter.backward(&loss).unwrap();
adapter.optimizer_step().unwrap();
let tmp_dir = std::env::temp_dir();
let checkpoint_path = tmp_dir.join("kan_test_ckpt");
let path_str = checkpoint_path.to_str().unwrap();
adapter.save_checkpoint(path_str).unwrap();
let mut adapter2 = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let metadata = adapter2.load_checkpoint(path_str).unwrap();
assert_eq!(metadata.model_type, "KAN");
assert_eq!(metadata.step, 1);
assert_eq!(adapter2.get_step(), 1);
// Clean up
let _ = std::fs::remove_file(format!("{}.json", path_str));
let _ = std::fs::remove_file(format!("{}.safetensors", path_str));
}
#[test]
fn test_validate() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let val_data: Vec<(Tensor, Tensor)> = (0..3)
.map(|_| {
let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap();
let target = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
(input, target)
})
.collect();
let val_loss = adapter.validate(&val_data).unwrap();
assert!(val_loss >= 0.0, "Validation loss should be non-negative");
assert!(adapter.latest_metrics.val_loss.is_some(), "val_loss should be set");
}
#[test]
fn test_validate_empty_errors() {
let cfg = make_config();
let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let result = adapter.validate(&[]);
assert!(result.is_err(), "Validating empty data should error");
}
}

1
ml/src/lib.rs
View File

@@ -692,6 +692,7 @@ pub mod evaluation; // DQN evaluation engine (backtest metrics, Sharpe ratio)
pub mod flash_attention;
pub mod hyperopt; // Bayesian hyperparameter optimization (egobox)
pub mod integration;
pub mod kan;
pub mod labeling;
pub mod liquid;
pub mod mamba;

2
ml/src/tlob/mod.rs
View File

@@ -7,6 +7,7 @@ pub mod analytics;
pub mod features;
pub mod mbp10_feature_extractor; // MBP-10 to TLOB feature extraction
pub mod performance;
pub mod trainable_adapter;
pub mod transformer;
// Re-export key types for external use
@@ -17,6 +18,7 @@ pub use features::{
TLOBFeatures as TLOBInputFeatures,
TLOB_FEATURE_COUNT,
};
pub use trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter};
pub use transformer::{TLOBConfig, TLOBMetrics, TLOBTransformer};
// Re-export transformer-specific TLOBFeatures with a different name to avoid conflicts

671
ml/src/tlob/trainable_adapter.rs Normal file
View File

@@ -0,0 +1,671 @@
//! UnifiedTrainable adapter for TLOB (Time Limit Order Book)
//!
//! Wraps a candle-based projection network to provide the UnifiedTrainable
//! interface for TLOB. The projection network maps flattened order book features
//! through hidden layers, enabling gradient-based training via the unified
//! training orchestrator.
//!
//! Architecture: input_linear(seq_len*feature_dim -> d_model) -> ReLU -> output_linear(d_model -> 1)
use candle_core::{backprop::GradStore, DType, Device, Module, Tensor};
use candle_nn::{linear, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::training::unified_trainer::{checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable};
use crate::MLError;
/// Configuration for the TLOB trainable adapter projection network.
///
/// This is separate from the ONNX-oriented `TLOBConfig` in `transformer.rs`.
/// It defines the architecture for the candle-based projection network used
/// by the unified training orchestrator.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TLOBAdapterConfig {
/// Model / hidden dimension for projection layers
pub d_model: usize,
/// Number of attention heads (reserved for future transformer layers)
pub num_heads: usize,
/// Number of projection layers (reserved for future deeper networks)
pub num_layers: usize,
/// Sequence length (number of time steps in LOB snapshot)
pub seq_len: usize,
/// Feature dimension per time step (e.g. 51 for TLOB features)
pub feature_dim: usize,
}
impl Default for TLOBAdapterConfig {
fn default() -> Self {
Self {
d_model: 128,
num_heads: 4,
num_layers: 2,
seq_len: 128,
feature_dim: 51,
}
}
}
/// Adapter wrapping TLOB with a candle-based projection network for unified training.
///
/// The projection network has two linear layers:
/// - `input_linear`: projects from `seq_len * feature_dim` to `d_model`
/// - `output_linear`: projects from `d_model` to 1 (scalar prediction)
///
/// Training uses AdamW with gradient tracking via `GradStore`.
pub struct TLOBTrainableAdapter {
/// TLOB adapter configuration
config: TLOBAdapterConfig,
/// Candle variable map holding learnable parameters
var_map: VarMap,
/// Input projection layer (seq_len*feature_dim -> d_model)
input_linear: Linear,
/// Output projection layer (d_model -> 1)
output_linear: Linear,
/// AdamW optimizer
optimizer: AdamW,
/// Gradient store from last backward pass (consumed by optimizer_step)
grads: Option<GradStore>,
/// Device (CPU or CUDA)
device: Device,
/// Current learning rate
learning_rate: f64,
/// Current training step
step: usize,
/// Latest training metrics
latest_metrics: TrainingMetrics,
/// Loss history for rolling average
loss_history: Vec<f64>,
/// Last computed gradient norm
last_grad_norm: f64,
}
impl std::fmt::Debug for TLOBTrainableAdapter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TLOBTrainableAdapter")
.field("config", &self.config)
.field("device", &format!("{:?}", self.device))
.field("learning_rate", &self.learning_rate)
.field("step", &self.step)
.field("loss_history_len", &self.loss_history.len())
.field("last_grad_norm", &self.last_grad_norm)
.finish_non_exhaustive()
}
}
impl TLOBTrainableAdapter {
/// Create a new TLOB trainable adapter with projection network.
///
/// # Arguments
/// * `config` - TLOB adapter configuration specifying dimensions
/// * `device` - Device to create tensors on (CPU or CUDA)
///
/// # Returns
/// Initialized adapter ready for training
pub fn new(config: TLOBAdapterConfig, device: &Device) -> Result<Self, MLError> {
let var_map = VarMap::new();
let vb = VarBuilder::from_varmap(&var_map, DType::F32, device);
let input_dim = config.seq_len * config.feature_dim;
let input_linear = linear(input_dim, config.d_model, vb.pp("input"))
.map_err(|e| MLError::ModelError(format!("Failed to create input linear: {}", e)))?;
let output_linear = linear(config.d_model, 1, vb.pp("output"))
.map_err(|e| MLError::ModelError(format!("Failed to create output linear: {}", e)))?;
let learning_rate = 1e-3;
let all_vars = var_map.all_vars();
let optimizer = AdamW::new(
all_vars,
ParamsAdamW {
lr: learning_rate,
beta1: 0.9,
beta2: 0.999,
eps: 1e-8,
weight_decay: 1e-4,
},
)
.map_err(|e| MLError::ModelError(format!("Failed to create AdamW optimizer: {}", e)))?;
Ok(Self {
config,
var_map,
input_linear,
output_linear,
optimizer,
grads: None,
device: device.clone(),
learning_rate,
step: 0,
latest_metrics: TrainingMetrics::default(),
loss_history: Vec::new(),
last_grad_norm: 0.0,
})
}
/// Access the underlying TLOB adapter configuration.
pub fn tlob_config(&self) -> &TLOBAdapterConfig {
&self.config
}
}
impl UnifiedTrainable for TLOBTrainableAdapter {
fn model_type(&self) -> &str {
"TLOB"
}
fn device(&self) -> &Device {
&self.device
}
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
// input shape: [batch, seq_len, feature_dim] or [batch, seq_len*feature_dim]
// Flatten to [batch, seq_len*feature_dim] if 3D
let dims = input.shape().dims();
let flat_input = if dims.len() == 3 {
let batch = dims.first().copied().ok_or_else(|| {
MLError::ModelError("Empty input dimensions".to_string())
})?;
let flat_dim = self.config.seq_len * self.config.feature_dim;
input.reshape(&[batch, flat_dim]).map_err(|e| {
MLError::ModelError(format!("Failed to reshape input to 2D: {}", e))
})?
} else {
input.clone()
};
// input_linear: seq_len*feature_dim -> d_model
let hidden = self.input_linear.forward(&flat_input).map_err(|e| {
MLError::ModelError(format!("Input linear forward failed: {}", e))
})?;
// ReLU activation
let activated = hidden.relu().map_err(|e| {
MLError::ModelError(format!("ReLU activation failed: {}", e))
})?;
// output_linear: d_model -> 1
let output = self.output_linear.forward(&activated).map_err(|e| {
MLError::ModelError(format!("Output linear forward failed: {}", e))
})?;
Ok(output)
}
fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor, MLError> {
// MSE loss: mean((predictions - targets)^2)
let diff = predictions.sub(targets).map_err(|e| {
MLError::ModelError(format!("Loss subtraction failed: {}", e))
})?;
let squared = diff.powf(2.0).map_err(|e| {
MLError::ModelError(format!("Loss squaring failed: {}", e))
})?;
let loss = squared.mean_all().map_err(|e| {
MLError::ModelError(format!("Loss mean failed: {}", e))
})?;
Ok(loss)
}
fn backward(&mut self, loss: &Tensor) -> Result<f64, MLError> {
let grads = loss.backward().map_err(|e| {
MLError::TrainingError(format!("Backward pass failed: {}", e))
})?;
// Compute gradient norm across all variables for monitoring
let mut grad_norm_sq = 0.0;
let vars_lock = self
.var_map
.data()
.lock()
.map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?;
for (_name, var) in vars_lock.iter() {
if let Some(grad) = grads.get(var.as_tensor()) {
let norm = grad
.sqr()
.and_then(|s| s.sum_all())
.and_then(|s| s.to_scalar::<f32>())
.map_err(|e| {
MLError::ModelError(format!("Failed to compute grad norm: {}", e))
})?;
grad_norm_sq += norm as f64;
}
}
drop(vars_lock);
let grad_norm = grad_norm_sq.sqrt();
self.last_grad_norm = grad_norm;
self.latest_metrics.grad_norm = Some(grad_norm);
self.grads = Some(grads);
Ok(grad_norm)
}
fn optimizer_step(&mut self) -> Result<(), MLError> {
if let Some(grads) = self.grads.take() {
self.optimizer.step(&grads).map_err(|e| {
MLError::TrainingError(format!("Optimizer step failed: {}", e))
})?;
}
self.step += 1;
Ok(())
}
fn zero_grad(&mut self) -> Result<(), MLError> {
// Gradients are implicitly zeroed by creating a new GradStore in backward().
// No explicit zeroing needed with candle's approach.
Ok(())
}
fn get_learning_rate(&self) -> f64 {
self.learning_rate
}
fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
self.learning_rate = lr;
// Recreate optimizer with new learning rate
let all_vars = self.var_map.all_vars();
self.optimizer = AdamW::new(
all_vars,
ParamsAdamW {
lr,
beta1: 0.9,
beta2: 0.999,
eps: 1e-8,
weight_decay: 1e-4,
},
)
.map_err(|e| {
MLError::ModelError(format!("Failed to recreate optimizer with new lr: {}", e))
})?;
Ok(())
}
fn get_step(&self) -> usize {
self.step
}
fn collect_metrics(&self) -> TrainingMetrics {
let mut metrics = self.latest_metrics.clone();
// Rolling average of recent losses
if !self.loss_history.is_empty() {
let recent: Vec<f64> = self.loss_history.iter().rev().take(100).copied().collect();
metrics.loss = recent.iter().sum::<f64>() / recent.len() as f64;
}
metrics.learning_rate = self.learning_rate;
// TLOB-specific custom metrics
metrics
.custom_metrics
.insert("training_steps".to_string(), self.step as f64);
metrics.custom_metrics.insert(
"loss_history_len".to_string(),
self.loss_history.len() as f64,
);
metrics.custom_metrics.insert(
"d_model".to_string(),
self.config.d_model as f64,
);
metrics.custom_metrics.insert(
"seq_len".to_string(),
self.config.seq_len as f64,
);
metrics.custom_metrics.insert(
"feature_dim".to_string(),
self.config.feature_dim as f64,
);
metrics.custom_metrics.insert(
"num_layers".to_string(),
self.config.num_layers as f64,
);
metrics
}
fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
let metadata = CheckpointMetadata {
model_type: "TLOB".to_string(),
version: "1.0.0".to_string(),
epoch: 0,
step: self.step,
timestamp: std::time::SystemTime::now(),
config: serde_json::to_value(&self.config).map_err(|e| {
MLError::SerializationError {
reason: format!("Failed to serialize config: {}", e),
}
})?,
metrics: self.collect_metrics(),
};
// Save JSON metadata
checkpoint::save_metadata(&metadata, checkpoint_path)?;
// Save model weights via safetensors
let safetensors_path = format!("{}.safetensors", checkpoint_path);
let vars_lock = self
.var_map
.data()
.lock()
.map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?;
let mut tensors: HashMap<String, Tensor> = HashMap::new();
for (name, var) in vars_lock.iter() {
tensors.insert(name.clone(), var.as_tensor().clone());
}
drop(vars_lock);
candle_core::safetensors::save(&tensors, &safetensors_path).map_err(|e| {
MLError::CheckpointError(format!("Failed to save safetensors: {}", e))
})?;
tracing::info!(
"Saved TLOB checkpoint to {} (step {})",
checkpoint_path,
self.step
);
Ok(checkpoint_path.to_string())
}
fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
let metadata = checkpoint::load_metadata(checkpoint_path)?;
if metadata.model_type != "TLOB" {
return Err(MLError::CheckpointError(format!(
"Invalid model type in checkpoint: expected TLOB, got {}",
metadata.model_type
)));
}
// Load weights from safetensors
let safetensors_path = format!("{}.safetensors", checkpoint_path);
let tensors = candle_core::safetensors::load(&safetensors_path, &self.device).map_err(
|e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)),
)?;
// Set loaded tensors into VarMap
let vars_lock = self
.var_map
.data()
.lock()
.map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?;
for (name, tensor) in &tensors {
if let Some(var) = vars_lock.get(name) {
var.set(tensor).map_err(|e| {
MLError::CheckpointError(format!("Failed to set var {}: {}", name, e))
})?;
} else {
tracing::warn!("Checkpoint contains unknown variable: {}", name);
}
}
drop(vars_lock);
// Restore training state
self.step = metadata.step;
self.latest_metrics = metadata.metrics.clone();
tracing::info!(
"Loaded TLOB checkpoint from {} (step {})",
checkpoint_path,
metadata.step
);
Ok(metadata)
}
fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result<f64, MLError> {
if val_data.is_empty() {
return Err(MLError::ValidationError {
message: "Empty validation dataset".to_string(),
});
}
let mut total_loss = 0.0;
let mut count = 0usize;
for (input, target) in val_data {
let prediction = self.forward(input)?;
let loss = self.compute_loss(&prediction, target)?;
let loss_val = loss.to_scalar::<f32>().map_err(|e| MLError::ValidationError {
message: format!("Failed to extract loss value: {}", e),
})?;
total_loss += loss_val as f64;
count += 1;
}
let avg_loss = if count > 0 {
total_loss / count as f64
} else {
0.0
};
self.latest_metrics.val_loss = Some(avg_loss);
tracing::debug!("TLOB validation loss: {:.6}", avg_loss);
Ok(avg_loss)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_config() -> TLOBAdapterConfig {
TLOBAdapterConfig {
d_model: 32,
num_heads: 2,
num_layers: 2,
seq_len: 32,
feature_dim: 51,
}
}
#[test]
fn test_model_type_returns_tlob() {
let cfg = make_config();
let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
assert_eq!(adapter.model_type(), "TLOB");
}
#[test]
fn test_device_returns_cpu() {
let cfg = make_config();
let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
assert!(matches!(adapter.device(), &Device::Cpu));
}
#[test]
fn test_forward_produces_output() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
// 3D input: [batch=2, seq_len=32, feature_dim=51]
let input =
Tensor::zeros(&[2, cfg.seq_len, cfg.feature_dim], DType::F32, &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let dims = output.shape().dims();
assert_eq!(dims.len(), 2);
assert_eq!(dims[0], 2); // batch
assert_eq!(dims[1], 1); // scalar prediction
}
#[test]
fn test_forward_accepts_2d_input() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
// 2D input: [batch=2, seq_len*feature_dim]
let flat_dim = cfg.seq_len * cfg.feature_dim;
let input = Tensor::zeros(&[2, flat_dim], DType::F32, &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let dims = output.shape().dims();
assert_eq!(dims.len(), 2);
assert_eq!(dims[0], 2);
assert_eq!(dims[1], 1);
}
#[test]
fn test_compute_loss_returns_scalar() {
let cfg = make_config();
let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let preds = Tensor::new(&[[1.0f32], [2.0]], &Device::Cpu).unwrap();
let targets = Tensor::new(&[[1.5f32], [2.5]], &Device::Cpu).unwrap();
let loss = adapter.compute_loss(&preds, &targets).unwrap();
let loss_val: f32 = loss.to_scalar().unwrap();
// MSE of (0.5^2 + 0.5^2)/2 = 0.25
assert!((loss_val - 0.25).abs() < 1e-5, "Expected ~0.25, got {}", loss_val);
}
#[test]
fn test_backward_returns_grad_norm() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
let flat_dim = cfg.seq_len * cfg.feature_dim;
let input = Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap();
let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let loss = adapter.compute_loss(&output, &targets).unwrap();
let grad_norm = adapter.backward(&loss).unwrap();
assert!(grad_norm >= 0.0, "Gradient norm should be non-negative");
assert!(adapter.grads.is_some(), "Grads should be stored after backward");
}
#[test]
fn test_train_step_cycle() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
assert_eq!(adapter.get_step(), 0);
// Full train cycle: zero_grad -> forward -> loss -> backward -> optimizer_step
adapter.zero_grad().unwrap();
let flat_dim = cfg.seq_len * cfg.feature_dim;
let input = Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap();
let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let loss = adapter.compute_loss(&output, &targets).unwrap();
let loss_val: f64 = loss.to_scalar::<f32>().unwrap() as f64;
adapter.backward(&loss).unwrap();
adapter.optimizer_step().unwrap();
assert!(adapter.get_step() >= 1);
// Record loss and check metrics
adapter.loss_history.push(loss_val);
let metrics = adapter.collect_metrics();
assert!(metrics.loss >= 0.0);
assert_eq!(
metrics.custom_metrics.get("training_steps").copied(),
Some(1.0)
);
}
#[test]
fn test_learning_rate_get_set() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let original_lr = adapter.get_learning_rate();
assert!((original_lr - 1e-3).abs() < 1e-10);
adapter.set_learning_rate(5e-4).unwrap();
assert!((adapter.get_learning_rate() - 5e-4).abs() < 1e-10);
}
#[test]
fn test_collect_metrics() {
let cfg = make_config();
let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let metrics = adapter.collect_metrics();
assert!(metrics.custom_metrics.contains_key("training_steps"));
assert!(metrics.custom_metrics.contains_key("d_model"));
assert!(metrics.custom_metrics.contains_key("seq_len"));
assert!(metrics.custom_metrics.contains_key("feature_dim"));
assert!(metrics.custom_metrics.contains_key("num_layers"));
assert_eq!(metrics.custom_metrics.get("d_model").copied(), Some(32.0));
assert_eq!(metrics.custom_metrics.get("seq_len").copied(), Some(32.0));
assert_eq!(metrics.custom_metrics.get("feature_dim").copied(), Some(51.0));
}
#[test]
fn test_checkpoint_roundtrip() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
// Run a training step to have non-zero state
let flat_dim = cfg.seq_len * cfg.feature_dim;
let input = Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap();
let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
let output = adapter.forward(&input).unwrap();
let loss = adapter.compute_loss(&output, &targets).unwrap();
adapter.backward(&loss).unwrap();
adapter.optimizer_step().unwrap();
// Save checkpoint
let tmp_dir = std::env::temp_dir();
let checkpoint_path = tmp_dir.join("tlob_test_ckpt");
let path_str = checkpoint_path.to_str().unwrap();
adapter.save_checkpoint(path_str).unwrap();
// Load into fresh adapter
let mut adapter2 = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let metadata = adapter2.load_checkpoint(path_str).unwrap();
assert_eq!(metadata.model_type, "TLOB");
assert_eq!(metadata.step, 1);
assert_eq!(adapter2.get_step(), 1);
// Clean up checkpoint files
let _ = std::fs::remove_file(format!("{}.json", path_str));
let _ = std::fs::remove_file(format!("{}.safetensors", path_str));
}
#[test]
fn test_validate_returns_loss() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap();
let flat_dim = cfg.seq_len * cfg.feature_dim;
let val_data: Vec<(Tensor, Tensor)> = (0..3)
.map(|_| {
let input =
Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap();
let target = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap();
(input, target)
})
.collect();
let val_loss = adapter.validate(&val_data).unwrap();
assert!(val_loss >= 0.0, "Validation loss should be non-negative");
assert!(
adapter.latest_metrics.val_loss.is_some(),
"val_loss should be set after validate"
);
}
#[test]
fn test_validate_empty_errors() {
let cfg = make_config();
let mut adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap();
let result = adapter.validate(&[]);
assert!(result.is_err(), "Validating empty data should error");
}
}