f946dcd952
WAVE 22: All examples, benchmarks, and data loaders updated Files Modified (41 files): - DQN examples: 7 files (train_dqn, evaluate_dqn, validate_dqn, etc.) - PPO examples: 6 files (train_ppo, continuous_ppo, benchmark_ppo, etc.) - TFT examples: 9 files (train_tft, validate_tft, benchmark_tft, etc.) - MAMBA-2 examples: 3 files (train_mamba2, verify_dimensions, etc.) - Benchmarks: 5 files (cuda_speedup, weight_caching, future_decoder, etc.) - Data loaders: 7 files (parquet_utils, dbn_sequence_loader, tlob_loader, etc.) - Integration: 4 files (load_parquet_data, streaming loaders, etc.) Key Changes: - state_dim: 225 → 54 (DQN, PPO) - input_dim: 225 → 54 (TFT) - d_model: 225 → 54 (MAMBA-2) - Memory: 1.8KB → 0.43KB per vector (76% reduction) - All tensor shapes updated: (batch, 225) → (batch, 54) Agents Deployed: 5 parallel agents Validation: cargo check PASSING Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com>
76 lines
2.8 KiB
Rust
76 lines
2.8 KiB
Rust
use candle_core::{Device, Tensor};
|
|
use ml::memory_optimization::quantization::Quantizer;
|
|
use ml::tft::{quantized_tft::QuantizedTemporalFusionTransformer, TFTConfig};
|
|
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("Testing forward_future_decoder implementation...\n");
|
|
|
|
// Create TFT config
|
|
let config = TFTConfig {
|
|
input_dim: 54,
|
|
hidden_dim: 256,
|
|
num_heads: 8,
|
|
num_known_features: 10,
|
|
prediction_horizon: 10,
|
|
..Default::default()
|
|
};
|
|
|
|
let device = Device::Cpu;
|
|
let qtft = QuantizedTemporalFusionTransformer::new_with_device(config, device.clone())?;
|
|
|
|
// Test 1: Create test future features [batch=2, horizon=10, features=10]
|
|
println!("Test 1: Basic forward pass");
|
|
let batch_size = 2;
|
|
let horizon = 10;
|
|
let num_features = 10;
|
|
|
|
let future_features = Tensor::randn(0f32, 1f32, (batch_size, horizon, num_features), &device)?;
|
|
|
|
println!(" Input shape: {:?}", future_features.dims());
|
|
|
|
// Create decoder weights [hidden_dim=256, num_features=10]
|
|
let weight_data: Vec<f32> = (0..256 * 10).map(|i| (i as f32 * 0.01).sin()).collect();
|
|
let weights_tensor = Tensor::from_slice(&weight_data, (256, 10), &device)?;
|
|
|
|
// Create quantizer and quantize the weights
|
|
let mut quantizer = ml::memory_optimization::quantization::Quantizer::new(
|
|
ml::memory_optimization::quantization::QuantizationConfig {
|
|
quant_type: ml::memory_optimization::quantization::QuantizationType::Int8,
|
|
per_channel: false,
|
|
symmetric: true,
|
|
calibration_samples: None,
|
|
},
|
|
device.clone(),
|
|
);
|
|
let quantized_weights = quantizer.quantize_tensor(&weights_tensor, "decoder")?;
|
|
|
|
// Run forward pass
|
|
let output = qtft.forward_future_decoder(&future_features, &quantized_weights)?;
|
|
|
|
println!(" Output shape: {:?}", output.dims());
|
|
println!(" Expected: [2, 10, 256]");
|
|
|
|
// Validate output shape
|
|
assert_eq!(output.dims(), &[2, 10, 256], "Output shape mismatch!");
|
|
println!(" ✓ Shape validation passed\n");
|
|
|
|
// Test 2: Check output is not all zeros
|
|
println!("Test 2: Output non-zero validation");
|
|
let output_sum = output.sum_all()?.to_vec0::<f32>()?;
|
|
println!(" Output sum: {}", output_sum);
|
|
assert!(output_sum.abs() > 1e-6, "Output should not be all zeros");
|
|
println!(" ✓ Non-zero validation passed\n");
|
|
|
|
// Test 3: Broadcasting correctness
|
|
println!("Test 3: Different batch sizes");
|
|
for batch in [1, 4, 8] {
|
|
let test_features = Tensor::randn(0f32, 1f32, (batch, 10, 10), &device)?;
|
|
let test_output = qtft.forward_future_decoder(&test_features, &quantized_weights)?;
|
|
assert_eq!(test_output.dims(), &[batch, 10, 256]);
|
|
println!(" ✓ Batch size {} works correctly", batch);
|
|
}
|
|
|
|
println!("\n✅ All tests passed!");
|
|
Ok(())
|
|
}
|