fix(tft): static context broadcast — unblocks 2 adapter tests

apply_static_context() tried gpu_add([1,9,32], [1,1,32]) — shape
mismatch. Static context is time-invariant, must broadcast across
sequence dimension before adding to temporal features.

Both TFT adapter tests now pass. Zero ignored in ensemble adapters.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 17:36:13 +01:00
parent 4776819666
commit c4f3b643e2
2 changed files with 53 additions and 7 deletions

View File

@@ -585,9 +585,57 @@ impl TemporalFusionTransformer {
temporal: &GpuTensor,
static_context: &GpuTensor,
) -> Result<GpuTensor, MLError> {
// 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)
}

View File

@@ -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)