daf771c38d
Every to_vec()/memcpy_dtoh across 48 files audited and annotated: - ~100 false positives: Rust slice .to_vec() (cpu-side, never touches GPU) - ~25 gpu-exit: legitimate scalar readbacks (loss, grad_norm, epoch state) - ~20 test-only readbacks: gated by #[cfg(test)] scope - ~10 cpu-side uploads: .to_vec() before from_vec() GPU upload - ~3 checkpoint exports: export_to_host at epoch boundary Annotations use inline comments: // cpu-side, // gpu-exit:, // test-only Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
585 lines
21 KiB
Rust
585 lines
21 KiB
Rust
//! Multi-timeframe LSTM encoder and feature fusion.
|
|
//!
|
|
//! Resamples 1-minute OHLCV bars into 5m / 15m / 1h bars, encodes each
|
|
//! timeframe independently with a single-layer LSTM, and fuses the four
|
|
//! hidden-state embeddings through a linear projection.
|
|
//!
|
|
//! Architecture:
|
|
//! ```text
|
|
//! 1m bars ──> [LSTM 6->64] ──> emb_1m (64) ─┐
|
|
//! ├─ resample 5m ──> [LSTM 6->64] ──> emb_5m (64) ─┤
|
|
//! ├─ resample 15m ──> [LSTM 6->64] ──> emb_15m (64) ─┼─> [Concat 256] ─> [Linear 128] ─> macro_state
|
|
//! └─ resample 1h ──> [LSTM 6->64] ──> emb_1h (64) ─┘
|
|
//! ```
|
|
|
|
use std::collections::VecDeque;
|
|
use std::sync::Arc;
|
|
|
|
use ml_core::cuda_autograd::GpuVarStore;
|
|
use ml_core::device::MlDevice;
|
|
use ml_core::native_types::NativeDevice;
|
|
|
|
use cudarc::driver::CudaStream;
|
|
use ml_supervised::gpu_tensor::{
|
|
GpuLinear, GpuTensor, gpu_sigmoid, gpu_tanh, gpu_mul, gpu_add,
|
|
gpu_cat_dim1, gpu_narrow_2d, gpu_matmul, gpu_transpose,
|
|
};
|
|
|
|
use super::bar_resampler::BarResampler;
|
|
use crate::types::OHLCVBar;
|
|
use crate::MLError;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Configuration
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Configuration for the multi-timeframe encoder.
|
|
#[derive(Debug, Clone)]
|
|
pub struct MultiTimeframeConfig {
|
|
/// Number of input features per bar (default: 6 = O/H/L/C/V/returns).
|
|
pub input_dim: usize,
|
|
/// Hidden dimension of each per-timeframe LSTM (default: 64).
|
|
pub hidden_dim: usize,
|
|
/// Output dimension after fusion projection (default: 128).
|
|
pub output_dim: usize,
|
|
/// Number of recent bars to keep per timeframe (default: 20).
|
|
pub history_len: usize,
|
|
}
|
|
|
|
impl Default for MultiTimeframeConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
input_dim: 6,
|
|
hidden_dim: 64,
|
|
output_dim: 128,
|
|
history_len: 20,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// LstmEncoder -- single-layer LSTM cell (manual implementation)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A single-layer LSTM encoder that processes a sequence and returns the
|
|
/// final hidden state as an embedding vector.
|
|
///
|
|
/// Uses cuBLAS-backed matmul via GpuTensor for gate computations.
|
|
///
|
|
/// LSTM equations:
|
|
/// ```text
|
|
/// gates = W_ih * x_t + b_ih + W_hh * h_{t-1} + b_hh
|
|
/// (i, f, g, o) = split(gates, 4)
|
|
/// i_t = sigmoid(i), f_t = sigmoid(f), g_t = tanh(g), o_t = sigmoid(o)
|
|
/// c_t = f_t * c_{t-1} + i_t * g_t
|
|
/// h_t = o_t * tanh(c_t)
|
|
/// ```
|
|
#[derive(Debug)]
|
|
pub struct LstmEncoder {
|
|
/// Input-to-hidden weights [4*hidden, input] — GpuTensor.
|
|
w_ih: GpuTensor,
|
|
/// Hidden-to-hidden weights [4*hidden, hidden] — GpuTensor.
|
|
w_hh: GpuTensor,
|
|
/// Input-to-hidden bias [4*hidden] — GpuTensor.
|
|
b_ih: GpuTensor,
|
|
/// Hidden-to-hidden bias [4*hidden] — GpuTensor.
|
|
b_hh: GpuTensor,
|
|
hidden_dim: usize,
|
|
stream: Arc<CudaStream>,
|
|
}
|
|
|
|
impl LstmEncoder {
|
|
/// Build a new LSTM encoder with Xavier-initialized weights on the CUDA stream.
|
|
pub fn new(input_dim: usize, hidden_dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
|
|
use rand::Rng;
|
|
let gate_dim = 4 * hidden_dim;
|
|
let limit_ih = (6.0 / (input_dim + hidden_dim) as f64).sqrt() as f32;
|
|
let limit_hh = (6.0 / (hidden_dim + hidden_dim) as f64).sqrt() as f32;
|
|
|
|
let mut rng = rand::thread_rng();
|
|
let w_ih_data: Vec<f32> = (0..gate_dim * input_dim).map(|_| rng.gen_range(-limit_ih..limit_ih)).collect();
|
|
let w_hh_data: Vec<f32> = (0..gate_dim * hidden_dim).map(|_| rng.gen_range(-limit_hh..limit_hh)).collect();
|
|
|
|
let w_ih = GpuTensor::from_vec(w_ih_data, &[gate_dim, input_dim], stream)?;
|
|
let w_hh = GpuTensor::from_vec(w_hh_data, &[gate_dim, hidden_dim], stream)?;
|
|
let b_ih = GpuTensor::zeros(&[1, gate_dim], stream)?;
|
|
let b_hh = GpuTensor::zeros(&[1, gate_dim], stream)?;
|
|
|
|
Ok(Self {
|
|
w_ih,
|
|
w_hh,
|
|
b_ih,
|
|
b_hh,
|
|
hidden_dim,
|
|
stream: Arc::clone(stream),
|
|
})
|
|
}
|
|
|
|
/// Run the LSTM over a sequence and return the final hidden state.
|
|
///
|
|
/// * `seq` -- GpuTensor of shape `(seq_len, input_dim)`
|
|
///
|
|
/// Returns a GpuTensor of shape `(1, hidden_dim)` (the final h).
|
|
pub fn forward(&self, seq: &GpuTensor) -> Result<GpuTensor, MLError> {
|
|
let seq_len = seq.dim(0)?;
|
|
|
|
let mut h = GpuTensor::zeros(&[1, self.hidden_dim], &self.stream)?;
|
|
let mut c = GpuTensor::zeros(&[1, self.hidden_dim], &self.stream)?;
|
|
|
|
// Transpose weight matrices once for efficient matmul: W^T
|
|
let w_ih_t = gpu_transpose(&self.w_ih)?;
|
|
let w_hh_t = gpu_transpose(&self.w_hh)?;
|
|
|
|
for t in 0..seq_len {
|
|
// x_t: (1, input_dim) -- narrow row t
|
|
let x_t = gpu_narrow_2d(seq, 0, t, 1)?;
|
|
|
|
// gates = x_t @ W_ih^T + h @ W_hh^T + b_ih + b_hh
|
|
let xw = gpu_matmul(&x_t, &w_ih_t)?;
|
|
let hw = gpu_matmul(&h, &w_hh_t)?;
|
|
let gates = gpu_add(&gpu_add(&gpu_add(&xw, &self.b_ih)?, &hw)?, &self.b_hh)?;
|
|
|
|
let hd = self.hidden_dim;
|
|
|
|
// Split gates into 4 chunks along dim 1
|
|
let i_gate = gpu_narrow_2d(&gates, 1, 0, hd)?;
|
|
let f_gate = gpu_narrow_2d(&gates, 1, hd, hd)?;
|
|
let g_gate = gpu_narrow_2d(&gates, 1, 2 * hd, hd)?;
|
|
let o_gate = gpu_narrow_2d(&gates, 1, 3 * hd, hd)?;
|
|
|
|
let i_sig = gpu_sigmoid(&i_gate)?;
|
|
let f_sig = gpu_sigmoid(&f_gate)?;
|
|
let g_tanh = gpu_tanh(&g_gate)?;
|
|
let o_sig = gpu_sigmoid(&o_gate)?;
|
|
|
|
// c_t = f_t * c_{t-1} + i_t * g_t
|
|
let fc = gpu_mul(&f_sig, &c)?;
|
|
let ig = gpu_mul(&i_sig, &g_tanh)?;
|
|
c = gpu_add(&fc, &ig)?;
|
|
|
|
// h_t = o_t * tanh(c_t)
|
|
let c_tanh = gpu_tanh(&c)?;
|
|
h = gpu_mul(&o_sig, &c_tanh)?;
|
|
}
|
|
|
|
Ok(h)
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// MultiTimeframeEncoder
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Encoder that processes OHLCV bars at four timeframes (1m, 5m, 15m, 1h)
|
|
/// and fuses the per-timeframe LSTM embeddings into a single macro-state vector.
|
|
///
|
|
/// All compute uses cuBLAS-backed GpuTensor. Output is converted to Candle Tensor
|
|
/// at the API boundary.
|
|
#[derive(Debug)]
|
|
pub struct MultiTimeframeEncoder {
|
|
lstm_1m: LstmEncoder,
|
|
lstm_5m: LstmEncoder,
|
|
lstm_15m: LstmEncoder,
|
|
lstm_1h: LstmEncoder,
|
|
projection: GpuLinear,
|
|
resampler: BarResampler,
|
|
history_1m: VecDeque<OHLCVBar>,
|
|
history_5m: VecDeque<OHLCVBar>,
|
|
history_15m: VecDeque<OHLCVBar>,
|
|
history_1h: VecDeque<OHLCVBar>,
|
|
config: MultiTimeframeConfig,
|
|
device: NativeDevice,
|
|
stream: Arc<CudaStream>,
|
|
}
|
|
|
|
impl MultiTimeframeEncoder {
|
|
/// Build a new encoder with GPU-native LSTM encoders and cuBLAS projection.
|
|
pub fn new(config: MultiTimeframeConfig, stream: &Arc<CudaStream>, device: &NativeDevice) -> Result<Self, MLError> {
|
|
let concat_dim = config.hidden_dim * 4; // 4 timeframes
|
|
|
|
let lstm_1m = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
|
|
let lstm_5m = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
|
|
let lstm_15m = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
|
|
let lstm_1h = LstmEncoder::new(config.input_dim, config.hidden_dim, stream)?;
|
|
|
|
let projection = GpuLinear::new(concat_dim, config.output_dim, stream)?;
|
|
|
|
let history_len = config.history_len;
|
|
|
|
Ok(Self {
|
|
lstm_1m,
|
|
lstm_5m,
|
|
lstm_15m,
|
|
lstm_1h,
|
|
projection,
|
|
resampler: BarResampler::new(),
|
|
history_1m: VecDeque::with_capacity(history_len),
|
|
history_5m: VecDeque::with_capacity(history_len),
|
|
history_15m: VecDeque::with_capacity(history_len),
|
|
history_1h: VecDeque::with_capacity(history_len),
|
|
config,
|
|
device: device.clone(),
|
|
stream: Arc::clone(stream),
|
|
})
|
|
}
|
|
|
|
/// Build with a fresh `GpuVarStore` on the specified device (convenience).
|
|
pub fn with_device(
|
|
config: MultiTimeframeConfig,
|
|
device: &NativeDevice,
|
|
) -> Result<(Self, GpuVarStore), MLError> {
|
|
let ordinal = device.cuda_ordinal().ok_or_else(|| {
|
|
MLError::ConfigError("MultiTimeframeEncoder requires CUDA device".to_owned())
|
|
})?;
|
|
let ml_dev = MlDevice::cuda(ordinal)?;
|
|
let stream = ml_dev.cuda_stream()?;
|
|
let vars = GpuVarStore::new(Arc::clone(stream));
|
|
let encoder = Self::new(config, stream, device)?;
|
|
Ok((encoder, vars))
|
|
}
|
|
|
|
/// Ingest a 1-minute bar, update internal ring buffers, and return
|
|
/// the fused macro-state embedding of shape `(1, output_dim)`.
|
|
///
|
|
/// The resampler converts the 1m bar into higher-timeframe bars when
|
|
/// enough constituent bars have been collected. Ring buffers are capped
|
|
/// at `history_len`.
|
|
pub fn push_bar(&mut self, bar: OHLCVBar) -> Result<GpuTensor, MLError> {
|
|
// Update resampler
|
|
let (bar_5m, bar_15m, bar_1h) = self.resampler.push(bar);
|
|
|
|
// Update ring buffers
|
|
push_ring(&mut self.history_1m, bar, self.config.history_len);
|
|
if let Some(b) = bar_5m {
|
|
push_ring(&mut self.history_5m, b, self.config.history_len);
|
|
}
|
|
if let Some(b) = bar_15m {
|
|
push_ring(&mut self.history_15m, b, self.config.history_len);
|
|
}
|
|
if let Some(b) = bar_1h {
|
|
push_ring(&mut self.history_1h, b, self.config.history_len);
|
|
}
|
|
|
|
self.encode()
|
|
}
|
|
|
|
/// Encode the current ring-buffer contents and return the fused state.
|
|
///
|
|
/// For any timeframe with no history yet, a zero embedding is used.
|
|
/// Returns a Candle Tensor at the API boundary.
|
|
pub fn encode(&self) -> Result<GpuTensor, MLError> {
|
|
let emb_1m = self.encode_timeframe(&self.lstm_1m, &self.history_1m)?;
|
|
let emb_5m = self.encode_timeframe(&self.lstm_5m, &self.history_5m)?;
|
|
let emb_15m = self.encode_timeframe(&self.lstm_15m, &self.history_15m)?;
|
|
let emb_1h = self.encode_timeframe(&self.lstm_1h, &self.history_1h)?;
|
|
|
|
// Concat along feature dim: (1, 4*hidden_dim) using GpuTensor ops
|
|
let cat_12 = gpu_cat_dim1(&emb_1m, &emb_5m)?;
|
|
let cat_123 = gpu_cat_dim1(&cat_12, &emb_15m)?;
|
|
let concat = gpu_cat_dim1(&cat_123, &emb_1h)?;
|
|
|
|
// Project to output_dim (cuBLAS sgemm)
|
|
self.projection.forward(&concat)
|
|
}
|
|
|
|
/// Reset the resampler and all history buffers.
|
|
pub fn reset(&mut self) {
|
|
self.resampler.reset();
|
|
self.history_1m.clear();
|
|
self.history_5m.clear();
|
|
self.history_15m.clear();
|
|
self.history_1h.clear();
|
|
}
|
|
|
|
/// Return a reference to the config.
|
|
pub fn config(&self) -> &MultiTimeframeConfig {
|
|
&self.config
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Internal helpers
|
|
// -----------------------------------------------------------------------
|
|
|
|
/// Encode a single timeframe's history through the given LSTM.
|
|
/// Returns `(1, hidden_dim)` GpuTensor.
|
|
fn encode_timeframe(
|
|
&self,
|
|
lstm: &LstmEncoder,
|
|
history: &VecDeque<OHLCVBar>,
|
|
) -> Result<GpuTensor, MLError> {
|
|
if history.is_empty() {
|
|
// No data yet -- return zeros
|
|
return GpuTensor::zeros(&[1, self.config.hidden_dim], &self.stream);
|
|
}
|
|
|
|
let seq = bars_to_gpu_tensor(history, &self.stream)?;
|
|
lstm.forward(&seq)
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Utility functions
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Convert an OHLCV bar to a 6-dim feature vector:
|
|
/// `[open, high, low, close, volume, returns]`.
|
|
///
|
|
/// `prev_close` is used to compute returns (`(close - prev_close) / prev_close`).
|
|
/// If `prev_close` is `None` or zero, returns is set to 0.0.
|
|
pub fn bar_to_features(bar: &OHLCVBar, prev_close: Option<f64>) -> [f64; 6] {
|
|
let returns = match prev_close {
|
|
Some(pc) if pc.abs() > f64::EPSILON => (bar.close - pc) / pc,
|
|
_ => 0.0,
|
|
};
|
|
[bar.open, bar.high, bar.low, bar.close, bar.volume, returns]
|
|
}
|
|
|
|
/// Convert a sequence of OHLCV bars to a `(seq_len, 6)` GpuTensor.
|
|
fn bars_to_gpu_tensor(bars: &VecDeque<OHLCVBar>, stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
|
|
let len = bars.len();
|
|
let mut data = Vec::with_capacity(len * 6);
|
|
|
|
let mut prev_close: Option<f64> = None;
|
|
for bar in bars {
|
|
let feats = bar_to_features(bar, prev_close);
|
|
for &f in &feats {
|
|
data.push(f as f32);
|
|
}
|
|
prev_close = Some(bar.close);
|
|
}
|
|
|
|
GpuTensor::from_vec(data, &[len, 6], stream)
|
|
}
|
|
|
|
/// Push a bar into a ring buffer, popping the oldest if at capacity.
|
|
fn push_ring(buf: &mut VecDeque<OHLCVBar>, bar: OHLCVBar, max_len: usize) {
|
|
if buf.len() >= max_len {
|
|
buf.pop_front();
|
|
}
|
|
buf.push_back(bar);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use chrono::{TimeZone, Utc};
|
|
use ml_core::cuda_autograd::stream_ops::{gpu_sub, gpu_abs};
|
|
|
|
fn make_bar(minute: u32, close: f64) -> OHLCVBar {
|
|
OHLCVBar {
|
|
timestamp: Utc
|
|
.with_ymd_and_hms(2026, 1, 1, 10, minute % 60, 0)
|
|
.single()
|
|
.unwrap_or_else(Utc::now),
|
|
open: close - 1.0,
|
|
high: close + 2.0,
|
|
low: close - 2.0,
|
|
close,
|
|
volume: 1000.0,
|
|
}
|
|
}
|
|
|
|
fn make_config() -> MultiTimeframeConfig {
|
|
MultiTimeframeConfig {
|
|
input_dim: 6,
|
|
hidden_dim: 64,
|
|
output_dim: 128,
|
|
history_len: 20,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_encoder_output_shape() {
|
|
let config = make_config();
|
|
let (mut encoder, _vars) =
|
|
MultiTimeframeEncoder::with_device(config.clone(), &NativeDevice::Cuda(0))
|
|
.expect("encoder creation should succeed");
|
|
|
|
// Push some bars
|
|
for i in 0..10 {
|
|
let result = encoder.push_bar(make_bar(i, 100.0 + f64::from(i)));
|
|
assert!(result.is_ok(), "push_bar should succeed: {:?}", result.err());
|
|
}
|
|
|
|
let output = encoder.encode().expect("encode should succeed");
|
|
let dims = &output.shape;
|
|
assert_eq!(dims.len(), 2, "output should be 2D");
|
|
assert_eq!(
|
|
dims.first().copied().unwrap_or(0),
|
|
1,
|
|
"batch dim should be 1"
|
|
);
|
|
assert_eq!(
|
|
dims.last().copied().unwrap_or(0),
|
|
config.output_dim,
|
|
"feature dim should be output_dim={}",
|
|
config.output_dim
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_encoder_deterministic() {
|
|
let config = make_config();
|
|
let (mut encoder, _vars) =
|
|
MultiTimeframeEncoder::with_device(config, &NativeDevice::Cuda(0))
|
|
.expect("encoder creation should succeed");
|
|
|
|
// Push bars
|
|
for i in 0..5 {
|
|
encoder
|
|
.push_bar(make_bar(i, 100.0 + f64::from(i)))
|
|
.expect("push should work");
|
|
}
|
|
|
|
let out1 = encoder.encode().expect("encode 1");
|
|
let out2 = encoder.encode().expect("encode 2");
|
|
|
|
// Same state -> same output: compute |out1 - out2| and sum on host
|
|
let diff_tensor = gpu_sub(&out1, &out2).expect("sub");
|
|
let abs_tensor = gpu_abs(&diff_tensor).expect("abs");
|
|
let abs_vec = abs_tensor.to_vec().expect("to_vec"); // test-only readback
|
|
let diff: f32 = abs_vec.iter().sum();
|
|
|
|
assert!(
|
|
diff < 1e-6,
|
|
"same input should produce same output, diff={}",
|
|
diff
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_config_defaults() {
|
|
let config = MultiTimeframeConfig::default();
|
|
assert_eq!(config.input_dim, 6);
|
|
assert_eq!(config.hidden_dim, 64);
|
|
assert_eq!(config.output_dim, 128);
|
|
assert_eq!(config.history_len, 20);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bar_to_features() {
|
|
let bar = OHLCVBar {
|
|
timestamp: Utc
|
|
.with_ymd_and_hms(2026, 1, 1, 10, 0, 0)
|
|
.single()
|
|
.unwrap_or_else(Utc::now),
|
|
open: 100.0,
|
|
high: 105.0,
|
|
low: 95.0,
|
|
close: 102.0,
|
|
volume: 5000.0,
|
|
};
|
|
|
|
// No previous close -> returns = 0
|
|
let feats = bar_to_features(&bar, None);
|
|
assert!((feats[0] - 100.0).abs() < f64::EPSILON, "open");
|
|
assert!((feats[1] - 105.0).abs() < f64::EPSILON, "high");
|
|
assert!((feats[2] - 95.0).abs() < f64::EPSILON, "low");
|
|
assert!((feats[3] - 102.0).abs() < f64::EPSILON, "close");
|
|
assert!((feats[4] - 5000.0).abs() < f64::EPSILON, "volume");
|
|
assert!((feats[5] - 0.0).abs() < f64::EPSILON, "returns with no prev");
|
|
|
|
// With previous close: returns = (102 - 100) / 100 = 0.02
|
|
let feats = bar_to_features(&bar, Some(100.0));
|
|
assert!(
|
|
(feats[5] - 0.02).abs() < 1e-10,
|
|
"returns should be 0.02, got {}",
|
|
feats[5]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bar_to_features_zero_prev_close() {
|
|
let bar = make_bar(0, 102.0);
|
|
let feats = bar_to_features(&bar, Some(0.0));
|
|
assert!(
|
|
feats[5].abs() < f64::EPSILON,
|
|
"returns should be 0.0 when prev_close is 0"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_empty_history_produces_output() {
|
|
// Even with no bars pushed, encode() should succeed (zero embeddings)
|
|
let config = make_config();
|
|
let (encoder, _vars) =
|
|
MultiTimeframeEncoder::with_device(config.clone(), &NativeDevice::Cuda(0))
|
|
.expect("encoder creation should succeed");
|
|
|
|
let output = encoder.encode().expect("encode on empty history should work");
|
|
let dims = &output.shape;
|
|
assert_eq!(dims.last().copied().unwrap_or(0), config.output_dim);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lstm_encoder_single_step() {
|
|
let ml_dev = MlDevice::cuda(0).expect("CUDA required");
|
|
let stream = std::sync::Arc::clone(ml_dev.cuda_stream().expect("stream"));
|
|
let lstm = LstmEncoder::new(6, 32, &stream).expect("lstm creation");
|
|
|
|
// Single timestep: (1, 6)
|
|
let input = GpuTensor::randn(&[1, 6], 1.0, &stream).expect("input tensor");
|
|
let out = lstm.forward(&input).expect("lstm forward");
|
|
assert_eq!(out.shape.len(), 2);
|
|
assert_eq!(out.dim(0).unwrap_or(0), 1);
|
|
assert_eq!(out.dim(1).unwrap_or(0), 32);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lstm_encoder_multi_step() {
|
|
let ml_dev = MlDevice::cuda(0).expect("CUDA required");
|
|
let stream = std::sync::Arc::clone(ml_dev.cuda_stream().expect("stream"));
|
|
let lstm = LstmEncoder::new(6, 64, &stream).expect("lstm creation");
|
|
|
|
// 10 timesteps: (10, 6)
|
|
let input = GpuTensor::randn(&[10, 6], 1.0, &stream).expect("input tensor");
|
|
let out = lstm.forward(&input).expect("lstm forward");
|
|
assert_eq!(out.dim(1).unwrap_or(0), 64);
|
|
}
|
|
|
|
#[test]
|
|
fn test_push_ring_eviction() {
|
|
let mut buf = VecDeque::new();
|
|
let max_len = 3;
|
|
|
|
for i in 0..5 {
|
|
push_ring(&mut buf, make_bar(i, 100.0 + f64::from(i)), max_len);
|
|
}
|
|
|
|
assert_eq!(buf.len(), 3, "ring buffer should cap at max_len");
|
|
// Should contain bars for minutes 2, 3, 4
|
|
let front = buf.front().expect("front exists");
|
|
assert!((front.close - 102.0).abs() < f64::EPSILON);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reset_clears_state() {
|
|
let config = make_config();
|
|
let (mut encoder, _vars) =
|
|
MultiTimeframeEncoder::with_device(config, &NativeDevice::Cuda(0))
|
|
.expect("encoder creation");
|
|
|
|
for i in 0..10 {
|
|
encoder
|
|
.push_bar(make_bar(i, 100.0 + f64::from(i)))
|
|
.expect("push");
|
|
}
|
|
|
|
encoder.reset();
|
|
|
|
// After reset, encode should return zero-based output
|
|
let output = encoder.encode().expect("encode after reset");
|
|
let abs_out = gpu_abs(&output).expect("abs");
|
|
let vals = abs_out.to_vec().expect("to_vec"); // test-only readback
|
|
let sum: f32 = vals.iter().sum();
|
|
|
|
// With all-zero embeddings going through the projection (which has bias),
|
|
// the output is just the projection bias. That's fine.
|
|
assert!(sum.is_finite(), "output should be finite after reset");
|
|
}
|
|
}
|