86 lines
3.3 KiB
Rust
86 lines
3.3 KiB
Rust
use rtx_tensor::{Tensor, Device};
|
|
use pinn_mre_helmholtz::{Config, ForwardWorkspace, LffnUNet1D};
|
|
use rtx_nn::Linear;
|
|
use rtx_nn::layers::Module;
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_linear_forward_out() {
|
|
let device = Device::Cuda(0);
|
|
|
|
// Create a simple linear layer
|
|
let in_features = 128;
|
|
let out_features = 64;
|
|
let batch_size = 200;
|
|
|
|
let linear = Linear::new(in_features, out_features, true, &device).expect("Failed to create linear");
|
|
|
|
// Create input tensor
|
|
let input = Tensor::randn(&[batch_size, in_features], &device).expect("Failed to create input");
|
|
println!("Input shape: {:?}", input.shape().dims());
|
|
|
|
// Create output buffer
|
|
let mut output = Tensor::zeros(&[batch_size, out_features], &device).expect("Failed to create output");
|
|
println!("Output buffer shape: {:?}", output.shape().dims());
|
|
|
|
// Test regular forward
|
|
println!("\nTesting regular forward...");
|
|
let result = linear.forward(&input).expect("Forward failed");
|
|
println!("Regular forward output shape: {:?}", result.shape().dims());
|
|
|
|
// Test forward_out
|
|
println!("\nTesting forward_out...");
|
|
linear.forward_out(&input, &mut output).expect("forward_out failed");
|
|
println!("forward_out completed!");
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_fused_then_linear() {
|
|
let device = Device::Cuda(0);
|
|
let n_points = 200;
|
|
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
// Create B weights
|
|
let b_weights = Tensor::randn(&[1, cfg.u_ff_dim], &device).expect("Failed to create B");
|
|
println!("B weights shape: {:?}", b_weights.shape().dims());
|
|
|
|
// Create input
|
|
let x_data: Vec<f32> = (0..n_points).map(|i| i as f32 / n_points as f32).collect();
|
|
let x_norm = Tensor::from_slice(&x_data, &[n_points, 1], &device).expect("Failed to create input");
|
|
println!("Input shape: {:?}", x_norm.shape().dims());
|
|
|
|
// Create features buffer
|
|
let mut features = Tensor::zeros(&[n_points, cfg.u_ff_dim * 2], &device).expect("Failed to create features");
|
|
println!("Features buffer shape: {:?}", features.shape().dims());
|
|
|
|
// Run fused fourier
|
|
println!("\nRunning fused_fourier_features_out...");
|
|
x_norm.fused_fourier_features_out(&b_weights, 2.0 * std::f32::consts::PI, &mut features)
|
|
.expect("Fused failed");
|
|
println!("Fused completed! Features shape: {:?}", features.shape().dims());
|
|
|
|
// Create linear layer (128 -> 64)
|
|
let linear = Linear::new(cfg.u_ff_dim * 2, cfg.u_hidden, true, &device).expect("Failed to create linear");
|
|
|
|
// Create output buffer
|
|
let mut output = Tensor::zeros(&[n_points, cfg.u_hidden], &device).expect("Failed to create output");
|
|
println!("Output buffer shape: {:?}", output.shape().dims());
|
|
|
|
// Test regular forward first
|
|
println!("\nTesting regular forward with fused output...");
|
|
let result = linear.forward(&features).expect("Forward failed");
|
|
println!("Regular forward completed! Shape: {:?}", result.shape().dims());
|
|
|
|
// Test forward_out
|
|
println!("\nTesting forward_out with fused output...");
|
|
linear.forward_out(&features, &mut output).expect("forward_out failed");
|
|
println!("forward_out completed!");
|
|
}
|