9c3d741a08
Move 17 library crates into crates/, CLI binary into bin/fxt, consolidate 10 test crates into testing/, split config crate from deployment config files. Root directory reduced from 38+ to ~17 directories. All Cargo.toml paths and build.rs proto refs updated. Co-Authored-By: Claude Opus 4.6 <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(())
|
|
}
|