gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: supervised GPU smoke test dimension mismatches

Browse Source 
- Liquid: remove seq_len from input (forward_loss uses 2D [batch, input_size])
- Mamba2: single sample input matching d_model=32
- xLSTM: remove seq_len from input (same as Liquid)
- Diffusion: create checkpoint dir for models using dir-based saves

All 8 supervised GPU smoke tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-21 09:14:53 +01:00
parent 65891816ee
commit c20d810067
1 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
crates/ml/tests/supervised_gpu_smoke_test.rs
View File

@@ -177,6 +177,8 @@ fn smoke_pipeline(
let tmp_dir = std::env::temp_dir().join(format!("gpu_smoke_{}", model_name.to_lowercase()));
std::fs::create_dir_all(&tmp_dir).unwrap();
let ckpt_path = tmp_dir.join("ckpt");
// Some models (Diffusion) treat checkpoint_path as a directory and write files inside it.
std::fs::create_dir_all(&ckpt_path).unwrap();
let save_result = adapter.save_checkpoint(ckpt_path.to_str().unwrap());
assert!(
save_result.is_ok(),
@@ -261,10 +263,11 @@ fn test_mamba2_gpu_smoke() {
adapter.device_name()
);
// Input: [batch=16, seq_len=8, d_model=32] flattened
let input = random_f32_data(16 * 8 * 32);
// Mamba2 compute_loss narrows to last step -> target is [batch=16, 1] flattened
let target = random_f32_data(16 * 1);
// Mamba2 forward_loss creates shape [1, input.len()]. The input_projection expects
// the last dimension to equal d_model. Output projection maps to 1.
// Input: [1, d_model=32] flattened, Target: [1, 1] flattened
let input = random_f32_data(32);
let target = random_f32_data(1);
smoke_pipeline(&mut adapter, &input, &target, "Mamba2");
}
@@ -364,9 +367,9 @@ fn test_liquid_gpu_smoke() {
adapter.device_name()
);
// Input: [batch=16, seq_len=4, input_size=10] flattened (3D)
let input = random_f32_data(16 * 4 * 10);
// Output: [batch=16, output_size=1] flattened
// Liquid forward_loss: batch = input.len() / input_size, so input is 2D flat [batch * input_size].
// Input: [batch=16, input_size=10] flattened, Output: [batch=16, output_size=1] flattened
let input = random_f32_data(16 * 10);
let target = random_f32_data(16 * 1);
smoke_pipeline(&mut adapter, &input, &target, "Liquid");
}
@@ -435,9 +438,9 @@ fn test_xlstm_gpu_smoke() {
adapter.device_name()
);
// Input: [batch=16, seq_len=4, input_dim=10] flattened (3D)
let input = random_f32_data(16 * 4 * 10);
// Output: [batch=16, 1] flattened
// xLSTM forward_loss: batch = input.len() / input_dim, so input is 2D flat [batch * input_dim].
// Input: [batch=16, input_dim=10] flattened, Output: [batch=16, output_dim=1] flattened
let input = random_f32_data(16 * 10);
let target = random_f32_data(16 * 1);
smoke_pipeline(&mut adapter, &input, &target, "xLSTM");
}