diff --git a/crates/ml-supervised/src/tft/mod.rs b/crates/ml-supervised/src/tft/mod.rs index 5b51ee603..5d34f9c37 100644 --- a/crates/ml-supervised/src/tft/mod.rs +++ b/crates/ml-supervised/src/tft/mod.rs @@ -585,9 +585,57 @@ impl TemporalFusionTransformer { temporal: &GpuTensor, static_context: &GpuTensor, ) -> Result { - // For 2D tensors: simple element-wise add (both should be [batch, hidden_dim]) - // Static context is already [batch, hidden_dim] from VSN+GRN. - // Temporal is [batch, hidden_dim] after the linear encoder. + // Both tensors may be 3D after variable selection + temporal processing. + // temporal: [batch, seq_len, hidden_dim] + // static_context: [batch, 1, hidden_dim] (VSN adds seq_len=1 for 2D static input) + // + // When shapes already match, add directly. Otherwise broadcast the + // static context along the sequence dimension so gpu_add succeeds. + if temporal.shape == static_context.shape { + return gpu_add(temporal, static_context); + } + + // 3D broadcast: tile static_context along dim-1 to match temporal's seq_len + if temporal.shape.len() == 3 && static_context.shape.len() == 3 { + let batch = temporal.dim(0)?; + let seq_len = temporal.dim(1)?; + let hidden = temporal.dim(2)?; + + let sc_batch = static_context.dim(0)?; + let sc_seq = static_context.dim(1)?; + let sc_hidden = static_context.dim(2)?; + + if sc_batch != batch || sc_hidden != hidden { + return Err(MLError::DimensionMismatch { + expected: batch * hidden, + actual: sc_batch * sc_hidden, + }); + } + + if sc_seq == 1 && seq_len > 1 { + // Broadcast: replicate the single-step static context across seq_len + let sc_host = static_context.to_vec()?; + let mut expanded = vec![0.0_f32; batch * seq_len * hidden]; + for b in 0..batch { + for s in 0..seq_len { + for h in 0..hidden { + let src = b * hidden + h; + let dst = b * seq_len * hidden + s * hidden + h; + if let (Some(&v), Some(slot)) = + (sc_host.get(src), expanded.get_mut(dst)) + { + *slot = v; + } + } + } + } + let expanded_tensor = + GpuTensor::from_vec(expanded, &[batch, seq_len, hidden], &temporal.stream)?; + return gpu_add(temporal, &expanded_tensor); + } + } + + // Fallback for 2D or matching shapes gpu_add(temporal, static_context) } diff --git a/crates/ml/src/ensemble/adapters/tft.rs b/crates/ml/src/ensemble/adapters/tft.rs index d5cdde6a1..8421e4897 100644 --- a/crates/ml/src/ensemble/adapters/tft.rs +++ b/crates/ml/src/ensemble/adapters/tft.rs @@ -382,8 +382,7 @@ mod tests { } #[test] - #[ignore = "TFT input dim mismatch (288 vs 32) — needs config alignment after CUDA migration"] - fn test_tft_adapter_buffers_and_predicts() { + fn test_tft_adapter_buffers_and_predicts() { let seq_len = 4; let adapter = TftInferenceAdapter::new(test_config(), seq_len) .expect("TftInferenceAdapter::new should succeed"); @@ -452,8 +451,7 @@ mod tests { } #[test] - #[ignore = "TFT input dim mismatch (288 vs 32) — needs config alignment after CUDA migration"] - fn test_tft_adapter_deterministic() { + fn test_tft_adapter_deterministic() { let seq_len = 4; // Create two adapters with identical configs let adapter1 = TftInferenceAdapter::new(test_config(), seq_len)