From c20d810067ea350cc59ff94166166c58ea1b98ce Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 21 Mar 2026 09:14:53 +0100 Subject: [PATCH] fix: supervised GPU smoke test dimension mismatches - 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) --- crates/ml/tests/supervised_gpu_smoke_test.rs | 23 +++++++++++--------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/crates/ml/tests/supervised_gpu_smoke_test.rs b/crates/ml/tests/supervised_gpu_smoke_test.rs index 7ade8dc99..e5d9aae65 100644 --- a/crates/ml/tests/supervised_gpu_smoke_test.rs +++ b/crates/ml/tests/supervised_gpu_smoke_test.rs @@ -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"); }