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:
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user