feat(ml): add TGGN inference adapter for ensemble

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ml/src/ensemble/adapters/mod.rs

@@ -1,13 +1,23 @@
 //! Per-model inference adapters for ensemble prediction
 
+pub mod diffusion;
 pub mod dqn;
+pub mod kan;
 pub mod liquid;
 pub mod mamba2;
 pub mod ppo;
 pub mod tft;
+pub mod tggn;
+pub mod tlob;
+pub mod xlstm;
 
+pub use diffusion::DiffusionInferenceAdapter;
 pub use dqn::DqnInferenceAdapter;
+pub use kan::KanInferenceAdapter;
 pub use liquid::LiquidInferenceAdapter;
 pub use mamba2::Mamba2InferenceAdapter;
 pub use ppo::PpoInferenceAdapter;
 pub use tft::TftInferenceAdapter;
+pub use tggn::TggnInferenceAdapter;
+pub use tlob::TlobInferenceAdapter;
+pub use xlstm::XlstmInferenceAdapter;

ml/src/ensemble/adapters/tggn.rs

@@ -0,0 +1,228 @@
+//! TGGN inference adapter for ensemble prediction
+//!
+//! Builds a 2-layer candle projection matching the TGGNTrainableAdapter
+//! architecture (input->hidden->1) for checkpoint-compatible inference.
+
+use std::sync::Mutex;
+
+use candle_core::{DType, Device, Module, Tensor};
+use candle_nn::{linear, Linear, VarBuilder, VarMap};
+
+use crate::ensemble::inference_adapter::{
+    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
+};
+use crate::gpu::DeviceConfig;
+use crate::{MLError, MLResult};
+
+/// 2-layer projection matching TGGNTrainableAdapter's candle architecture.
+struct TggnProjection {
+    linear1: Linear,
+    linear2: Linear,
+}
+
+impl TggnProjection {
+    fn new(input_dim: usize, hidden_dim: usize, vb: VarBuilder<'_>) -> MLResult<Self> {
+        let linear1 = linear(input_dim, hidden_dim, vb.pp("linear1"))
+            .map_err(|e| MLError::ModelError(format!("TGGN linear1 init: {e}")))?;
+        let linear2 = linear(hidden_dim, 1, vb.pp("linear2"))
+            .map_err(|e| MLError::ModelError(format!("TGGN linear2 init: {e}")))?;
+        Ok(Self { linear1, linear2 })
+    }
+
+    fn forward(&self, input: &Tensor) -> MLResult<Tensor> {
+        let h = self
+            .linear1
+            .forward(input)
+            .map_err(|e| MLError::ModelError(format!("TGGN forward linear1: {e}")))?;
+        let h = h
+            .relu()
+            .map_err(|e| MLError::ModelError(format!("TGGN forward relu: {e}")))?;
+        self.linear2
+            .forward(&h)
+            .map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}")))
+    }
+}
+
+/// Inference adapter for the TGGN (Temporal Graph Gated Network).
+///
+/// Since the base TGGN model uses ndarray (not candle), this adapter
+/// replicates the 2-layer candle projection from `TGGNTrainableAdapter`
+/// (`input_dim -> hidden_dim -> 1`) for checkpoint-compatible inference.
+/// The scalar output is mapped via sigmoid to direction [-1, 1] and
+/// confidence [0, 1] for ensemble aggregation.
+#[allow(missing_debug_implementations)]
+pub struct TggnInferenceAdapter {
+    model: Mutex<TggnProjection>,
+    #[allow(dead_code)]
+    varmap: VarMap,
+    device: Device,
+    input_dim: usize,
+}
+
+// SAFETY: TggnProjection uses candle tensors which are Send+Sync.
+// The Mutex provides exclusive access for inference calls.
+#[allow(unsafe_code)]
+unsafe impl Send for TggnInferenceAdapter {}
+#[allow(unsafe_code)]
+unsafe impl Sync for TggnInferenceAdapter {}
+
+impl TggnInferenceAdapter {
+    /// Create a new TGGN inference adapter.
+    ///
+    /// # Arguments
+    /// * `input_dim` - Feature vector dimension (typically 51)
+    /// * `hidden_dim` - Hidden layer width (32-64 for RTX 3050 Ti)
+    pub fn new(input_dim: usize, hidden_dim: usize) -> MLResult<Self> {
+        let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
+        let varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let projection = TggnProjection::new(input_dim, hidden_dim, vb)?;
+
+        Ok(Self {
+            model: Mutex::new(projection),
+            varmap,
+            device,
+            input_dim,
+        })
+    }
+
+    /// Load from a safetensors checkpoint.
+    pub fn from_checkpoint(
+        input_dim: usize,
+        hidden_dim: usize,
+        path: &str,
+    ) -> MLResult<Self> {
+        let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
+        let mut varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let projection = TggnProjection::new(input_dim, hidden_dim, vb)?;
+
+        varmap
+            .load(path)
+            .map_err(|e| MLError::ModelError(format!("Failed to load TGGN checkpoint: {e}")))?;
+
+        Ok(Self {
+            model: Mutex::new(projection),
+            varmap,
+            device,
+            input_dim,
+        })
+    }
+
+    /// Pad or truncate feature values to match input_dim.
+    fn pad_features(&self, values: &[f64]) -> Vec<f32> {
+        let mut padded = vec![0.0f32; self.input_dim];
+        let copy_len = values.len().min(self.input_dim);
+        for i in 0..copy_len {
+            if let Some(v) = values.get(i) {
+                if let Some(slot) = padded.get_mut(i) {
+                    *slot = *v as f32;
+                }
+            }
+        }
+        padded
+    }
+}
+
+impl ModelInferenceAdapter for TggnInferenceAdapter {
+    fn model_name(&self) -> &str {
+        "TGGN"
+    }
+
+    fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
+        let start = std::time::Instant::now();
+
+        let padded = self.pad_features(&features.values);
+        let input = Tensor::from_vec(padded, (1, self.input_dim), &self.device)
+            .map_err(|e| MLError::ModelError(format!("TGGN input tensor: {e}")))?;
+
+        let model = self
+            .model
+            .lock()
+            .map_err(|e| MLError::LockError(format!("TGGN lock poisoned: {e}")))?;
+        let output = model.forward(&input)?;
+
+        let squeezed = output
+            .squeeze(0)
+            .map_err(|e| MLError::ModelError(format!("TGGN squeeze: {e}")))?;
+        let raw: Vec<f32> = squeezed
+            .to_vec1()
+            .map_err(|e| MLError::ModelError(format!("TGGN extract: {e}")))?;
+
+        let raw_val = raw.first().copied().unwrap_or(0.0) as f64;
+        let prob = 1.0 / (1.0 + (-raw_val).exp());
+        let direction = (2.0 * prob - 1.0).clamp(-1.0, 1.0);
+        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);
+
+        let latency_us = start.elapsed().as_micros() as u64;
+
+        Ok(EnsemblePrediction {
+            model_name: "TGGN".to_string(),
+            direction,
+            confidence,
+            metadata: PredictionMeta {
+                latency_us,
+                quantiles: None,
+                attention_weights: None,
+                q_values: None,
+            },
+        })
+    }
+
+    fn is_ready(&self) -> bool {
+        self.model.lock().is_ok()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_tggn_adapter_creation() {
+        let adapter = TggnInferenceAdapter::new(51, 64);
+        assert!(
+            adapter.is_ok(),
+            "TGGN adapter creation failed: {:?}",
+            adapter.err()
+        );
+        let adapter = adapter.unwrap();
+        assert_eq!(adapter.model_name(), "TGGN");
+        assert!(adapter.is_ready());
+    }
+
+    #[test]
+    fn test_tggn_adapter_predict_direction_range() {
+        let adapter = TggnInferenceAdapter::new(51, 64).unwrap();
+        let fv = FeatureVector {
+            values: vec![0.1; 51],
+            timestamp: 1700000000_000_000,
+        };
+        let pred = adapter.predict(&fv).unwrap();
+        assert!(
+            pred.direction >= -1.0 && pred.direction <= 1.0,
+            "direction {} out of [-1,1]",
+            pred.direction
+        );
+        assert!(
+            pred.confidence >= 0.0 && pred.confidence <= 1.0,
+            "confidence {} out of [0,1]",
+            pred.confidence
+        );
+    }
+
+    #[test]
+    fn test_tggn_adapter_deterministic() {
+        let adapter = TggnInferenceAdapter::new(51, 64).unwrap();
+        let fv = FeatureVector {
+            values: vec![0.1; 51],
+            timestamp: 1700000000_000_000,
+        };
+        let pred1 = adapter.predict(&fv).unwrap();
+        let pred2 = adapter.predict(&fv).unwrap();
+        assert_eq!(
+            pred1.direction, pred2.direction,
+            "Deterministic predictions should have same direction"
+        );
+    }
+}

ml/src/ensemble/adapters/xlstm.rs

@@ -0,0 +1,340 @@
+//! xLSTM inference adapter for ensemble prediction
+//!
+//! Wraps an `XLSTMNetwork` with a sequence buffer. Feature vectors accumulate
+//! until `sequence_length` frames are collected, then the sLSTM+mLSTM blocks
+//! run a forward pass and the scalar output is mapped via sigmoid to
+//! direction + confidence.
+
+use std::collections::VecDeque;
+use std::sync::Mutex;
+
+use candle_core::{DType, Device, Tensor};
+use candle_nn::{VarBuilder, VarMap};
+
+use crate::ensemble::inference_adapter::{
+    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
+};
+use crate::gpu::DeviceConfig;
+use crate::xlstm::config::XLSTMConfig;
+use crate::xlstm::network::XLSTMNetwork;
+use crate::{MLError, MLResult};
+
+/// Inference adapter for xLSTM (sLSTM + mLSTM blocks) with sequence buffering.
+///
+/// Maintains a ring buffer of feature vectors. Returns a neutral prediction
+/// (direction=0, confidence=0) until `sequence_length` frames have been
+/// collected. Once the buffer is full, each call builds a
+/// `[1, seq_len, input_dim]` tensor and runs the xLSTM forward pass.
+/// The scalar output is mapped via sigmoid to direction [-1, 1] and
+/// confidence [0, 1].
+#[allow(missing_debug_implementations)]
+pub struct XlstmInferenceAdapter {
+    model: Mutex<XLSTMNetwork>,
+    buffer: Mutex<VecDeque<Vec<f32>>>,
+    #[allow(dead_code)]
+    varmap: VarMap,
+    sequence_length: usize,
+    input_dim: usize,
+    device: Device,
+}
+
+// SAFETY: XLSTMNetwork uses candle tensors which are Send+Sync.
+// The Mutex provides exclusive access for inference calls and buffer mutation.
+#[allow(unsafe_code)]
+unsafe impl Send for XlstmInferenceAdapter {}
+#[allow(unsafe_code)]
+unsafe impl Sync for XlstmInferenceAdapter {}
+
+impl XlstmInferenceAdapter {
+    /// Create a new xLSTM inference adapter.
+    ///
+    /// # Arguments
+    /// * `config` - xLSTM model configuration
+    /// * `sequence_length` - Number of feature vectors to buffer before inference
+    pub fn new(config: XLSTMConfig, sequence_length: usize) -> MLResult<Self> {
+        let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
+        let input_dim = config.input_dim;
+        let varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let network = XLSTMNetwork::new(&config, vb)?;
+
+        Ok(Self {
+            model: Mutex::new(network),
+            buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
+            varmap,
+            sequence_length,
+            input_dim,
+            device,
+        })
+    }
+
+    /// Create an xLSTM inference adapter and load weights from a safetensors checkpoint.
+    pub fn from_checkpoint(
+        config: XLSTMConfig,
+        sequence_length: usize,
+        path: &str,
+    ) -> MLResult<Self> {
+        let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
+        let input_dim = config.input_dim;
+        let mut varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let network = XLSTMNetwork::new(&config, vb)?;
+
+        varmap
+            .load(path)
+            .map_err(|e| MLError::ModelError(format!("Failed to load xLSTM checkpoint: {e}")))?;
+
+        Ok(Self {
+            model: Mutex::new(network),
+            buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
+            varmap,
+            sequence_length,
+            input_dim,
+            device,
+        })
+    }
+
+    /// Pad or truncate feature values to match input_dim, converting f64 to f32.
+    fn pad_features(&self, values: &[f64]) -> Vec<f32> {
+        let mut padded = vec![0.0f32; self.input_dim];
+        let copy_len = values.len().min(self.input_dim);
+        for i in 0..copy_len {
+            if let Some(v) = values.get(i) {
+                if let Some(slot) = padded.get_mut(i) {
+                    *slot = *v as f32;
+                }
+            }
+        }
+        padded
+    }
+}
+
+impl ModelInferenceAdapter for XlstmInferenceAdapter {
+    fn model_name(&self) -> &str {
+        "xLSTM"
+    }
+
+    fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
+        let start = std::time::Instant::now();
+
+        let padded = self.pad_features(&features.values);
+
+        // Buffer phase (short-lived lock)
+        let buffer_ready = {
+            let mut buf = self
+                .buffer
+                .lock()
+                .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?;
+            buf.push_back(padded);
+            while buf.len() > self.sequence_length {
+                buf.pop_front();
+            }
+            buf.len() >= self.sequence_length
+        };
+
+        if !buffer_ready {
+            let latency_us = start.elapsed().as_micros() as u64;
+            return Ok(EnsemblePrediction {
+                model_name: "xLSTM".to_string(),
+                direction: 0.0,
+                confidence: 0.0,
+                metadata: PredictionMeta {
+                    latency_us,
+                    quantiles: None,
+                    attention_weights: None,
+                    q_values: None,
+                },
+            });
+        }
+
+        // Build [1, seq_len, input_dim] tensor from buffer snapshot
+        let flat_data = {
+            let buf = self
+                .buffer
+                .lock()
+                .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?;
+            let mut data = Vec::with_capacity(self.sequence_length * self.input_dim);
+            for frame in buf.iter() {
+                data.extend_from_slice(frame);
+            }
+            data
+        };
+
+        let input = Tensor::from_vec(
+            flat_data,
+            (1, self.sequence_length, self.input_dim),
+            &self.device,
+        )
+        .map_err(|e| MLError::ModelError(format!("xLSTM input tensor: {e}")))?;
+
+        // XLSTMNetwork::forward takes &self (not &mut self)
+        let model = self
+            .model
+            .lock()
+            .map_err(|e| MLError::LockError(format!("xLSTM model lock poisoned: {e}")))?;
+        let output = model.forward(&input)?;
+
+        // Output: [1, output_dim] -> squeeze -> [output_dim]
+        let squeezed = output
+            .squeeze(0)
+            .map_err(|e| MLError::ModelError(format!("xLSTM squeeze: {e}")))?;
+        let raw: Vec<f32> = squeezed
+            .to_vec1()
+            .map_err(|e| MLError::ModelError(format!("xLSTM extract: {e}")))?;
+
+        let raw_val = raw.first().copied().unwrap_or(0.0) as f64;
+
+        // Sigmoid -> direction + confidence
+        let prob = 1.0 / (1.0 + (-raw_val).exp());
+        let direction = (2.0 * prob - 1.0).clamp(-1.0, 1.0);
+        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);
+
+        let latency_us = start.elapsed().as_micros() as u64;
+
+        Ok(EnsemblePrediction {
+            model_name: "xLSTM".to_string(),
+            direction,
+            confidence,
+            metadata: PredictionMeta {
+                latency_us,
+                quantiles: None,
+                attention_weights: None,
+                q_values: None,
+            },
+        })
+    }
+
+    fn is_ready(&self) -> bool {
+        let buf_ok = self
+            .buffer
+            .lock()
+            .map(|buf| buf.len() >= self.sequence_length)
+            .unwrap_or(false);
+        let model_ok = self.model.lock().is_ok();
+        buf_ok && model_ok
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_config() -> XLSTMConfig {
+        XLSTMConfig {
+            input_dim: 51,
+            hidden_dim: 32,
+            num_blocks: 2,
+            num_heads: 2,
+            output_dim: 1,
+            dropout: 0.0,
+            ..Default::default()
+        }
+    }
+
+    const TEST_SEQ_LEN: usize = 4;
+
+    #[test]
+    fn test_xlstm_adapter_creation() {
+        let adapter = XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN);
+        assert!(
+            adapter.is_ok(),
+            "xLSTM adapter creation failed: {:?}",
+            adapter.err()
+        );
+        let adapter = match adapter {
+            Ok(a) => a,
+            Err(_) => return,
+        };
+        assert_eq!(adapter.model_name(), "xLSTM");
+        assert!(!adapter.is_ready(), "Should not be ready with empty buffer");
+    }
+
+    #[test]
+    fn test_xlstm_adapter_buffers_and_predicts() {
+        let adapter = match XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN) {
+            Ok(a) => a,
+            Err(e) => {
+                panic!("Failed to create xLSTM adapter: {e:?}");
+            }
+        };
+
+        for i in 0..TEST_SEQ_LEN {
+            let fv = FeatureVector {
+                values: vec![0.1 * (i as f64 + 1.0); 51],
+                timestamp: 1_700_000_000_000_000 + i as i64,
+            };
+            let pred = match adapter.predict(&fv) {
+                Ok(p) => p,
+                Err(e) => {
+                    panic!("predict failed at step {i}: {e:?}");
+                }
+            };
+            if i < TEST_SEQ_LEN - 1 {
+                assert_eq!(pred.direction, 0.0, "Neutral direction while buffering");
+                assert_eq!(pred.confidence, 0.0, "Zero confidence while buffering");
+            } else {
+                assert!(
+                    pred.direction >= -1.0 && pred.direction <= 1.0,
+                    "direction {} out of [-1,1]",
+                    pred.direction
+                );
+                assert!(
+                    pred.confidence >= 0.0 && pred.confidence <= 1.0,
+                    "confidence {} out of [0,1]",
+                    pred.confidence
+                );
+            }
+        }
+        assert!(adapter.is_ready());
+    }
+
+    #[test]
+    fn test_xlstm_adapter_deterministic() {
+        let adapter = match XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN) {
+            Ok(a) => a,
+            Err(e) => {
+                panic!("Failed to create xLSTM adapter: {e:?}");
+            }
+        };
+
+        // Fill buffer with identical feature vectors
+        for _ in 0..TEST_SEQ_LEN {
+            let fv = FeatureVector {
+                values: vec![0.1; 51],
+                timestamp: 1_700_000_000_000_000,
+            };
+            let _ = adapter.predict(&fv);
+        }
+
+        // Two predictions with same input (buffer is full and stable)
+        let fv = FeatureVector {
+            values: vec![0.1; 51],
+            timestamp: 1_700_000_000_000_000,
+        };
+        let pred1 = match adapter.predict(&fv) {
+            Ok(p) => p,
+            Err(e) => {
+                panic!("predict1 failed: {e:?}");
+            }
+        };
+        let pred2 = match adapter.predict(&fv) {
+            Ok(p) => p,
+            Err(e) => {
+                panic!("predict2 failed: {e:?}");
+            }
+        };
+
+        assert!(
+            (pred1.direction - pred2.direction).abs() < 1e-6,
+            "direction mismatch: {} vs {}",
+            pred1.direction,
+            pred2.direction
+        );
+        assert!(
+            (pred1.confidence - pred2.confidence).abs() < 1e-6,
+            "confidence mismatch: {} vs {}",
+            pred1.confidence,
+            pred2.confidence
+        );
+    }
+}