701 lines
24 KiB
Rust
701 lines
24 KiB
Rust
//! Integration tests for weight loading.
|
|
//!
|
|
//! These tests require the generated weight files from convert_fno_weights.py.
|
|
|
|
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
|
use rtx_neural_operator::{FNO2d, SafeTensorsFile, load_config, load_fno2d_weights};
|
|
use rtx_nn::generic::GenericModule4D;
|
|
use rtx_tensor::generic::GenericTensor;
|
|
use std::path::Path;
|
|
|
|
const WEIGHTS_PATH: &str = "weights/fno/fno2d_darcy/model.safetensors";
|
|
const CONFIG_PATH: &str = "weights/fno/fno2d_darcy/config.json";
|
|
|
|
fn get_workspace_root() -> std::path::PathBuf {
|
|
// Navigate from crate directory to workspace root
|
|
let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
|
|
Path::new(&manifest_dir)
|
|
.parent()
|
|
.unwrap()
|
|
.parent()
|
|
.unwrap()
|
|
.parent()
|
|
.unwrap()
|
|
.to_path_buf()
|
|
}
|
|
|
|
#[test]
|
|
fn test_load_safetensors_file() {
|
|
let workspace = get_workspace_root();
|
|
let weights_path = workspace.join(WEIGHTS_PATH);
|
|
|
|
if !weights_path.exists() {
|
|
eprintln!(
|
|
"Skipping test: weights file not found at {:?}",
|
|
weights_path
|
|
);
|
|
eprintln!(
|
|
"Run: python scripts/convert_fno_weights.py --model fno2d_darcy --train --output weights/fno"
|
|
);
|
|
return;
|
|
}
|
|
|
|
let safetensors =
|
|
SafeTensorsFile::load(&weights_path).expect("Failed to load SafeTensors file");
|
|
|
|
// Check that expected tensors are present
|
|
let tensor_names = safetensors.tensor_names();
|
|
println!("Found {} tensors", tensor_names.len());
|
|
for name in &tensor_names {
|
|
println!(" {}", name);
|
|
}
|
|
|
|
// Check for either legacy or neuraloperator v2.0 format
|
|
let has_legacy = safetensors.contains("lifting.weight");
|
|
let has_v2 = safetensors.contains("lifting.fcs.0.weight");
|
|
assert!(
|
|
has_legacy || has_v2,
|
|
"Missing lifting weights (neither legacy nor v2.0 format)"
|
|
);
|
|
|
|
assert!(
|
|
safetensors.contains("spectral_conv.0.weights1_real"),
|
|
"Missing spectral_conv.0.weights1_real"
|
|
);
|
|
assert!(
|
|
safetensors.contains("projection.1.weight"),
|
|
"Missing projection.1.weight"
|
|
);
|
|
|
|
// Check lifting weight shape based on format
|
|
if has_v2 {
|
|
// neuraloperator v2.0 format: fc1 is [2*width, in_channels+2]
|
|
let fc1 = safetensors.get("lifting.fcs.0.weight").unwrap();
|
|
let fc1_shape = fc1.shape();
|
|
println!("Lifting fc1 shape: {:?}", fc1_shape);
|
|
// fc1 output should be 2*width = 64
|
|
assert_eq!(fc1_shape[0], 64, "Wrong fc1 output dimension");
|
|
// fc1 input should be in_channels + 2 = 3 (for darcy, data_channels=1 + 2 pos encoding)
|
|
assert_eq!(fc1_shape[1], 3, "Wrong fc1 input dimension");
|
|
|
|
let fc2 = safetensors.get("lifting.fcs.1.weight").unwrap();
|
|
let fc2_shape = fc2.shape();
|
|
println!("Lifting fc2 shape: {:?}", fc2_shape);
|
|
// fc2 should be [width, 2*width]
|
|
assert_eq!(fc2_shape[0], 32, "Wrong fc2 output dimension");
|
|
assert_eq!(fc2_shape[1], 64, "Wrong fc2 input dimension");
|
|
} else {
|
|
// Legacy format
|
|
let lifting = safetensors.get("lifting.weight").unwrap();
|
|
assert_eq!(lifting.shape(), &[32, 3], "Wrong lifting weight shape");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_load_config() {
|
|
let workspace = get_workspace_root();
|
|
let config_path = workspace.join(CONFIG_PATH);
|
|
|
|
if !config_path.exists() {
|
|
eprintln!("Skipping test: config file not found at {:?}", config_path);
|
|
return;
|
|
}
|
|
|
|
let config = load_config(&config_path).expect("Failed to load config");
|
|
|
|
assert_eq!(config.in_channels, 3);
|
|
assert_eq!(config.out_channels, 1);
|
|
assert_eq!(config.width, 32);
|
|
assert_eq!(config.n_modes, (12, 12));
|
|
assert_eq!(config.n_layers, 4);
|
|
assert_eq!(config.pde_type, "darcy");
|
|
}
|
|
|
|
#[test]
|
|
fn test_load_fno2d_weights() {
|
|
let workspace = get_workspace_root();
|
|
let weights_path = workspace.join(WEIGHTS_PATH);
|
|
|
|
if !weights_path.exists() {
|
|
eprintln!(
|
|
"Skipping test: weights file not found at {:?}",
|
|
weights_path
|
|
);
|
|
return;
|
|
}
|
|
|
|
// Check if weights are in neuraloperator v2.0 format
|
|
let safetensors = SafeTensorsFile::load(&weights_path).expect("Failed to load SafeTensors");
|
|
let is_v2_format = safetensors.contains("lifting.fcs.0.weight");
|
|
|
|
if !is_v2_format {
|
|
eprintln!("Skipping test: weights are in legacy format, not neuraloperator v2.0 format");
|
|
eprintln!("To generate v2.0 weights, run:");
|
|
eprintln!(
|
|
" python3 scripts/convert_fno_weights.py --model fno2d_darcy --train --output weights/fno"
|
|
);
|
|
return;
|
|
}
|
|
|
|
let weights = load_fno2d_weights(&weights_path).expect("Failed to load FNO2d weights");
|
|
|
|
// Check config was inferred correctly
|
|
// Note: in_channels is the data channels (without positional encoding)
|
|
// For Darcy, this is 1 (the coefficient field)
|
|
println!(
|
|
"Config: in_channels={}, out_channels={}, width={}, n_layers={}",
|
|
weights.config.in_channels,
|
|
weights.config.out_channels,
|
|
weights.config.width,
|
|
weights.config.n_layers
|
|
);
|
|
|
|
assert_eq!(weights.config.out_channels, 1);
|
|
assert_eq!(weights.config.width, 32);
|
|
assert_eq!(weights.config.n_layers, 4);
|
|
|
|
// Check lifting MLP weight shapes (neuraloperator v2.0 format)
|
|
// fc1: [2*width, in_channels+2] = [64, in_channels+2]
|
|
// fc2: [width, 2*width] = [32, 64]
|
|
let hidden_dim = 2 * weights.config.width; // 64
|
|
let in_with_pos = weights.config.in_channels + 2;
|
|
|
|
assert_eq!(
|
|
weights.lifting_fc1_weight.len(),
|
|
hidden_dim * in_with_pos,
|
|
"fc1 weight wrong size: got {}, expected {}",
|
|
weights.lifting_fc1_weight.len(),
|
|
hidden_dim * in_with_pos
|
|
);
|
|
assert_eq!(
|
|
weights.lifting_fc1_bias.len(),
|
|
hidden_dim,
|
|
"fc1 bias wrong size"
|
|
);
|
|
assert_eq!(
|
|
weights.lifting_fc2_weight.len(),
|
|
weights.config.width * hidden_dim,
|
|
"fc2 weight wrong size: got {}, expected {}",
|
|
weights.lifting_fc2_weight.len(),
|
|
weights.config.width * hidden_dim
|
|
);
|
|
assert_eq!(
|
|
weights.lifting_fc2_bias.len(),
|
|
weights.config.width,
|
|
"fc2 bias wrong size"
|
|
);
|
|
|
|
// Check other layers
|
|
assert_eq!(weights.spectral_weights.len(), 4);
|
|
assert_eq!(weights.conv_weights.len(), 4);
|
|
assert_eq!(weights.projection_weights.len(), 2);
|
|
|
|
// Check spectral conv weights
|
|
let expected_spectral_size = 32 * 32 * 12 * 12;
|
|
for (i, sw) in weights.spectral_weights.iter().enumerate() {
|
|
assert_eq!(
|
|
sw.weights1_real.len(),
|
|
expected_spectral_size,
|
|
"Layer {} weights1_real wrong size",
|
|
i
|
|
);
|
|
assert_eq!(
|
|
sw.weights1_imag.len(),
|
|
expected_spectral_size,
|
|
"Layer {} weights1_imag wrong size",
|
|
i
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno2d_from_weights_and_inference() {
|
|
let workspace = get_workspace_root();
|
|
let weights_path = workspace.join(WEIGHTS_PATH);
|
|
|
|
if !weights_path.exists() {
|
|
eprintln!(
|
|
"Skipping test: weights file not found at {:?}",
|
|
weights_path
|
|
);
|
|
return;
|
|
}
|
|
|
|
// Check if weights are in neuraloperator v2.0 format
|
|
let safetensors = SafeTensorsFile::load(&weights_path).expect("Failed to load SafeTensors");
|
|
let is_v2_format = safetensors.contains("lifting.fcs.0.weight");
|
|
|
|
if !is_v2_format {
|
|
eprintln!("Skipping test: weights are in legacy format, not neuraloperator v2.0 format");
|
|
eprintln!("To generate v2.0 weights, run:");
|
|
eprintln!(
|
|
" python3 scripts/convert_fno_weights.py --model fno2d_darcy --train --output weights/fno"
|
|
);
|
|
return;
|
|
}
|
|
|
|
let weights = load_fno2d_weights(&weights_path).expect("Failed to load FNO2d weights");
|
|
let device = CpuDevice::new();
|
|
|
|
let fno = FNO2d::<CpuBackend>::from_weights(&weights, &device)
|
|
.expect("Failed to create FNO2d from weights");
|
|
|
|
// Create input tensor using in_channels from config
|
|
// The model expects [batch, in_channels, height, width]
|
|
// Positional encoding is added internally, so we use in_channels (not in_channels+2)
|
|
let in_channels = weights.config.in_channels;
|
|
let input = GenericTensor::randn([1, in_channels, 64, 64], &device);
|
|
|
|
// Run inference
|
|
let output = fno.forward_4d(&input);
|
|
|
|
// Check output shape
|
|
assert_eq!(output.shape()[0], 1, "Wrong batch size");
|
|
assert_eq!(
|
|
output.shape()[1],
|
|
weights.config.out_channels,
|
|
"Wrong output channels"
|
|
);
|
|
assert_eq!(output.shape()[2], 64, "Wrong output height");
|
|
assert_eq!(output.shape()[3], 64, "Wrong output width");
|
|
|
|
// Check output is not all zeros
|
|
let output_data = output.to_vec();
|
|
let non_zero = output_data.iter().filter(|&&x| x.abs() > 1e-10).count();
|
|
assert!(non_zero > 0, "Output is all zeros");
|
|
|
|
println!("FNO2d inference successful!");
|
|
println!(" Input shape: {:?}", input.shape());
|
|
println!(" Output shape: {:?}", output.shape());
|
|
println!(
|
|
" Output range: [{:.4}, {:.4}]",
|
|
output_data.iter().cloned().reduce(f32::min).unwrap(),
|
|
output_data.iter().cloned().reduce(f32::max).unwrap()
|
|
);
|
|
}
|
|
|
|
const REFERENCE_INPUT_PATH: &str = "weights/fno/fno2d_darcy/reference_input.safetensors";
|
|
const REFERENCE_OUTPUT_PATH: &str = "weights/fno/fno2d_darcy/reference_output.safetensors";
|
|
|
|
use rtx_neural_operator::{LiftingMLP, SpectralConv2d};
|
|
|
|
/// Test that Rust FNO2d produces the same output as Python implementation.
|
|
///
|
|
/// This test loads reference input/output generated by Python and verifies
|
|
/// numerical equivalence within a tolerance.
|
|
#[test]
|
|
fn test_python_rust_numerical_equivalence() {
|
|
let workspace = get_workspace_root();
|
|
let weights_path = workspace.join(WEIGHTS_PATH);
|
|
let input_path = workspace.join(REFERENCE_INPUT_PATH);
|
|
let output_path = workspace.join(REFERENCE_OUTPUT_PATH);
|
|
|
|
// Check all required files exist
|
|
if !weights_path.exists() {
|
|
eprintln!(
|
|
"Skipping test: weights file not found at {:?}",
|
|
weights_path
|
|
);
|
|
eprintln!("Run: python scripts/generate_reference_outputs.py");
|
|
return;
|
|
}
|
|
if !input_path.exists() {
|
|
eprintln!(
|
|
"Skipping test: reference input not found at {:?}",
|
|
input_path
|
|
);
|
|
eprintln!("Run: python scripts/generate_reference_outputs.py");
|
|
return;
|
|
}
|
|
if !output_path.exists() {
|
|
eprintln!(
|
|
"Skipping test: reference output not found at {:?}",
|
|
output_path
|
|
);
|
|
eprintln!("Run: python scripts/generate_reference_outputs.py");
|
|
return;
|
|
}
|
|
|
|
// Check v2.0 format
|
|
let safetensors = SafeTensorsFile::load(&weights_path).expect("Failed to load weights");
|
|
if !safetensors.contains("lifting.fcs.0.weight") {
|
|
eprintln!("Skipping test: weights are in legacy format");
|
|
return;
|
|
}
|
|
|
|
// Load weights and create model
|
|
let weights = load_fno2d_weights(&weights_path).expect("Failed to load FNO2d weights");
|
|
let device = CpuDevice::new();
|
|
let fno = FNO2d::<CpuBackend>::from_weights(&weights, &device)
|
|
.expect("Failed to create FNO2d from weights");
|
|
|
|
// Load reference input
|
|
let input_safetensors =
|
|
SafeTensorsFile::load(&input_path).expect("Failed to load reference input");
|
|
let input_tensor_info = input_safetensors
|
|
.get("input")
|
|
.expect("Reference input missing 'input' tensor");
|
|
let input_data = input_tensor_info
|
|
.to_f32()
|
|
.expect("Failed to convert input to f32");
|
|
let input_shape = input_tensor_info.shape();
|
|
|
|
println!("Reference input shape: {:?}", input_shape);
|
|
|
|
// Create input tensor from reference data
|
|
let input = GenericTensor::from_slice(
|
|
&input_data,
|
|
[
|
|
input_shape[0],
|
|
input_shape[1],
|
|
input_shape[2],
|
|
input_shape[3],
|
|
],
|
|
&device,
|
|
);
|
|
|
|
// Run Rust inference
|
|
let rust_output = fno.forward_4d(&input);
|
|
let rust_output_data = rust_output.to_vec();
|
|
|
|
// Load reference output
|
|
let output_safetensors =
|
|
SafeTensorsFile::load(&output_path).expect("Failed to load reference output");
|
|
let output_tensor_info = output_safetensors
|
|
.get("output")
|
|
.expect("Reference output missing 'output' tensor");
|
|
let python_output_data = output_tensor_info
|
|
.to_f32()
|
|
.expect("Failed to convert output to f32");
|
|
let python_output_shape = output_tensor_info.shape();
|
|
|
|
println!("Reference output shape: {:?}", python_output_shape);
|
|
println!("Rust output shape: {:?}", rust_output.shape());
|
|
|
|
// Check shapes match
|
|
assert_eq!(
|
|
rust_output.shape()[0],
|
|
python_output_shape[0],
|
|
"Batch size mismatch"
|
|
);
|
|
assert_eq!(
|
|
rust_output.shape()[1],
|
|
python_output_shape[1],
|
|
"Channel count mismatch"
|
|
);
|
|
assert_eq!(
|
|
rust_output.shape()[2],
|
|
python_output_shape[2],
|
|
"Height mismatch"
|
|
);
|
|
assert_eq!(
|
|
rust_output.shape()[3],
|
|
python_output_shape[3],
|
|
"Width mismatch"
|
|
);
|
|
|
|
// Compare outputs
|
|
let n_elements = rust_output_data.len();
|
|
assert_eq!(n_elements, python_output_data.len(), "Output size mismatch");
|
|
|
|
// Calculate statistics
|
|
let mut max_abs_diff: f32 = 0.0;
|
|
let mut sum_abs_diff: f32 = 0.0;
|
|
let mut sum_sq_diff: f32 = 0.0;
|
|
let mut max_rel_diff: f32 = 0.0;
|
|
|
|
for (i, (rust_val, python_val)) in rust_output_data
|
|
.iter()
|
|
.zip(python_output_data.iter())
|
|
.enumerate()
|
|
{
|
|
let abs_diff = (rust_val - python_val).abs();
|
|
sum_abs_diff += abs_diff;
|
|
sum_sq_diff += abs_diff * abs_diff;
|
|
|
|
if abs_diff > max_abs_diff {
|
|
max_abs_diff = abs_diff;
|
|
}
|
|
|
|
// Relative difference (avoid division by zero)
|
|
let rel_diff = if python_val.abs() > 1e-10 {
|
|
abs_diff / python_val.abs()
|
|
} else {
|
|
abs_diff
|
|
};
|
|
if rel_diff > max_rel_diff {
|
|
max_rel_diff = rel_diff;
|
|
}
|
|
|
|
// Print first few differences for debugging
|
|
if i < 5 {
|
|
println!(" [{i}] rust={rust_val:.6}, python={python_val:.6}, diff={abs_diff:.6e}");
|
|
}
|
|
}
|
|
|
|
let mean_abs_diff = sum_abs_diff / n_elements as f32;
|
|
let rmse = (sum_sq_diff / n_elements as f32).sqrt();
|
|
|
|
println!("\nNumerical comparison:");
|
|
println!(" Max absolute difference: {:.6e}", max_abs_diff);
|
|
println!(" Mean absolute difference: {:.6e}", mean_abs_diff);
|
|
println!(" RMSE: {:.6e}", rmse);
|
|
println!(" Max relative difference: {:.6e}", max_rel_diff);
|
|
|
|
// Python output statistics
|
|
let python_mean: f32 = python_output_data.iter().sum::<f32>() / n_elements as f32;
|
|
let python_min = python_output_data.iter().cloned().reduce(f32::min).unwrap();
|
|
let python_max = python_output_data.iter().cloned().reduce(f32::max).unwrap();
|
|
|
|
println!(
|
|
"\nPython output: mean={:.6}, range=[{:.6}, {:.6}]",
|
|
python_mean, python_min, python_max
|
|
);
|
|
|
|
// Rust output statistics
|
|
let rust_mean: f32 = rust_output_data.iter().sum::<f32>() / n_elements as f32;
|
|
let rust_min = rust_output_data.iter().cloned().reduce(f32::min).unwrap();
|
|
let rust_max = rust_output_data.iter().cloned().reduce(f32::max).unwrap();
|
|
|
|
println!(
|
|
"Rust output: mean={:.6}, range=[{:.6}, {:.6}]",
|
|
rust_mean, rust_min, rust_max
|
|
);
|
|
|
|
// Assert numerical equivalence within tolerance
|
|
// FFT operations can accumulate small numerical differences, so we use relaxed tolerance
|
|
let tolerance = 1e-4;
|
|
assert!(
|
|
max_abs_diff < tolerance,
|
|
"Max absolute difference {:.6e} exceeds tolerance {:.6e}",
|
|
max_abs_diff,
|
|
tolerance
|
|
);
|
|
|
|
println!("\nPython vs Rust comparison PASSED!");
|
|
println!(
|
|
"Maximum difference {:.6e} is within tolerance {:.6e}",
|
|
max_abs_diff, tolerance
|
|
);
|
|
}
|
|
|
|
/// Test isolated spectral convolution layer against Python reference.
|
|
#[test]
|
|
fn test_spectral_conv_isolated() {
|
|
// Load debug data from Python script
|
|
let debug_path = std::path::Path::new("/tmp/spectral_conv_debug.safetensors");
|
|
if !debug_path.exists() {
|
|
eprintln!("Skipping test: debug data not found at {:?}", debug_path);
|
|
eprintln!("Run: python scripts/debug_spectral_conv.py");
|
|
return;
|
|
}
|
|
|
|
let safetensors = SafeTensorsFile::load(debug_path).expect("Failed to load debug data");
|
|
|
|
// Load tensors
|
|
let input_info = safetensors.get("input").expect("Missing input");
|
|
let w1_real_info = safetensors.get("w1_real").expect("Missing w1_real");
|
|
let w1_imag_info = safetensors.get("w1_imag").expect("Missing w1_imag");
|
|
let w2_real_info = safetensors.get("w2_real").expect("Missing w2_real");
|
|
let w2_imag_info = safetensors.get("w2_imag").expect("Missing w2_imag");
|
|
let expected_info = safetensors.get("output").expect("Missing output");
|
|
|
|
let input_data = input_info.to_f32().expect("Failed to convert input");
|
|
let w1_real = w1_real_info.to_f32().expect("Failed to convert w1_real");
|
|
let w1_imag = w1_imag_info.to_f32().expect("Failed to convert w1_imag");
|
|
let w2_real = w2_real_info.to_f32().expect("Failed to convert w2_real");
|
|
let w2_imag = w2_imag_info.to_f32().expect("Failed to convert w2_imag");
|
|
let expected_data = expected_info.to_f32().expect("Failed to convert output");
|
|
|
|
let input_shape = input_info.shape();
|
|
let w1_shape = w1_real_info.shape();
|
|
|
|
println!("Input shape: {:?}", input_shape);
|
|
println!("W1 shape: {:?}", w1_shape);
|
|
println!("Expected output shape: {:?}", expected_info.shape());
|
|
|
|
let batch = input_shape[0];
|
|
let in_ch = input_shape[1];
|
|
let height = input_shape[2];
|
|
let width = input_shape[3];
|
|
let out_ch = w1_shape[1];
|
|
let modes_h = w1_shape[2];
|
|
let modes_w = w1_shape[3];
|
|
|
|
println!(
|
|
"batch={}, in_ch={}, out_ch={}, height={}, width={}",
|
|
batch, in_ch, out_ch, height, width
|
|
);
|
|
println!("modes_h={}, modes_w={}", modes_h, modes_w);
|
|
|
|
// Create spectral conv layer with the debug weights
|
|
let device = CpuDevice::new();
|
|
let spectral_conv = SpectralConv2d::<CpuBackend>::from_weights_dual(
|
|
&w1_real, &w1_imag, &w2_real, &w2_imag, in_ch, out_ch, modes_h, modes_w, &device,
|
|
);
|
|
|
|
// Create input tensor
|
|
let input = GenericTensor::from_slice(&input_data, [batch, in_ch, height, width], &device);
|
|
|
|
// Run spectral conv
|
|
let output = spectral_conv.forward_4d(&input);
|
|
let output_data = output.to_vec();
|
|
|
|
println!(
|
|
"\nPython output: range=[{:.6}, {:.6}], first 4: {:?}",
|
|
expected_data.iter().cloned().reduce(f32::min).unwrap(),
|
|
expected_data.iter().cloned().reduce(f32::max).unwrap(),
|
|
&expected_data[..4]
|
|
);
|
|
println!(
|
|
"Rust output: range=[{:.6}, {:.6}], first 4: {:?}",
|
|
output_data.iter().cloned().reduce(f32::min).unwrap(),
|
|
output_data.iter().cloned().reduce(f32::max).unwrap(),
|
|
&output_data[..4]
|
|
);
|
|
|
|
// Calculate difference
|
|
let mut max_diff: f32 = 0.0;
|
|
for (r, p) in output_data.iter().zip(expected_data.iter()) {
|
|
let diff = (r - p).abs();
|
|
if diff > max_diff {
|
|
max_diff = diff;
|
|
}
|
|
}
|
|
|
|
println!("\nMax absolute difference: {:.6e}", max_diff);
|
|
|
|
// Check if outputs have similar magnitude (within 10x)
|
|
let python_max = expected_data
|
|
.iter()
|
|
.map(|x| x.abs())
|
|
.reduce(f32::max)
|
|
.unwrap();
|
|
let rust_max = output_data
|
|
.iter()
|
|
.map(|x| x.abs())
|
|
.reduce(f32::max)
|
|
.unwrap();
|
|
let ratio = rust_max / python_max;
|
|
println!("Magnitude ratio (Rust/Python): {:.4}", ratio);
|
|
}
|
|
|
|
/// Test isolated lifting MLP against Python reference.
|
|
#[test]
|
|
fn test_lifting_mlp_isolated() {
|
|
// Load debug data from Python script
|
|
let debug_path = std::path::Path::new("/tmp/lifting_mlp_debug.safetensors");
|
|
if !debug_path.exists() {
|
|
eprintln!("Skipping test: debug data not found at {:?}", debug_path);
|
|
eprintln!("Run: python scripts/debug_lifting_mlp.py");
|
|
return;
|
|
}
|
|
|
|
let safetensors = SafeTensorsFile::load(debug_path).expect("Failed to load debug data");
|
|
|
|
// Load tensors
|
|
let input_info = safetensors.get("input").expect("Missing input");
|
|
let fc1_weight_info = safetensors.get("fc1_weight").expect("Missing fc1_weight");
|
|
let fc1_bias_info = safetensors.get("fc1_bias").expect("Missing fc1_bias");
|
|
let fc2_weight_info = safetensors.get("fc2_weight").expect("Missing fc2_weight");
|
|
let fc2_bias_info = safetensors.get("fc2_bias").expect("Missing fc2_bias");
|
|
let expected_info = safetensors.get("output").expect("Missing output");
|
|
|
|
let input_data = input_info.to_f32().expect("Failed to convert input");
|
|
let fc1_weight = fc1_weight_info
|
|
.to_f32()
|
|
.expect("Failed to convert fc1_weight");
|
|
let fc1_bias = fc1_bias_info.to_f32().expect("Failed to convert fc1_bias");
|
|
let fc2_weight = fc2_weight_info
|
|
.to_f32()
|
|
.expect("Failed to convert fc2_weight");
|
|
let fc2_bias = fc2_bias_info.to_f32().expect("Failed to convert fc2_bias");
|
|
let expected_data = expected_info.to_f32().expect("Failed to convert output");
|
|
|
|
let input_shape = input_info.shape();
|
|
let fc1_shape = fc1_weight_info.shape();
|
|
let fc2_shape = fc2_weight_info.shape();
|
|
|
|
println!("Input shape: {:?}", input_shape);
|
|
println!("FC1 weight shape: {:?}", fc1_shape);
|
|
println!("FC2 weight shape: {:?}", fc2_shape);
|
|
println!("Expected output shape: {:?}", expected_info.shape());
|
|
|
|
let batch = input_shape[0];
|
|
let in_channels = input_shape[1];
|
|
let height = input_shape[2];
|
|
let width = input_shape[3];
|
|
let hidden_dim = fc1_shape[0];
|
|
let out_channels = fc2_shape[0];
|
|
|
|
println!(
|
|
"in_channels={}, hidden_dim={}, out_channels={}",
|
|
in_channels, hidden_dim, out_channels
|
|
);
|
|
|
|
// Create lifting MLP with the debug weights
|
|
let device = CpuDevice::new();
|
|
let lifting = LiftingMLP::<CpuBackend>::from_weights(
|
|
&fc1_weight,
|
|
&fc1_bias,
|
|
&fc2_weight,
|
|
&fc2_bias,
|
|
in_channels,
|
|
out_channels,
|
|
&device,
|
|
);
|
|
|
|
// Create input tensor
|
|
let input =
|
|
GenericTensor::from_slice(&input_data, [batch, in_channels, height, width], &device);
|
|
|
|
// Run lifting MLP
|
|
let output = lifting.forward_4d(&input);
|
|
let output_data = output.to_vec();
|
|
|
|
println!(
|
|
"\nPython output: range=[{:.4}, {:.4}], first 4: {:?}",
|
|
expected_data.iter().cloned().reduce(f32::min).unwrap(),
|
|
expected_data.iter().cloned().reduce(f32::max).unwrap(),
|
|
&expected_data[..4]
|
|
);
|
|
println!(
|
|
"Rust output: range=[{:.4}, {:.4}], first 4: {:?}",
|
|
output_data.iter().cloned().reduce(f32::min).unwrap(),
|
|
output_data.iter().cloned().reduce(f32::max).unwrap(),
|
|
&output_data[..4]
|
|
);
|
|
|
|
// Calculate difference
|
|
let mut max_diff: f32 = 0.0;
|
|
for (r, p) in output_data.iter().zip(expected_data.iter()) {
|
|
let diff: f32 = (r - p).abs();
|
|
if diff > max_diff {
|
|
max_diff = diff;
|
|
}
|
|
}
|
|
|
|
println!("\nMax absolute difference: {:.6e}", max_diff);
|
|
|
|
// Check magnitude ratio
|
|
let python_max: f32 = expected_data
|
|
.iter()
|
|
.map(|x| x.abs())
|
|
.reduce(f32::max)
|
|
.unwrap();
|
|
let rust_max: f32 = output_data
|
|
.iter()
|
|
.map(|x| x.abs())
|
|
.reduce(f32::max)
|
|
.unwrap();
|
|
let ratio = rust_max / python_max;
|
|
println!("Magnitude ratio (Rust/Python): {:.4}", ratio);
|
|
|
|
// GELU approximation can have small differences
|
|
assert!(
|
|
max_diff < 1e-3,
|
|
"Max difference {} exceeds tolerance",
|
|
max_diff
|
|
);
|
|
}
|