385 lines
12 KiB
Rust
385 lines
12 KiB
Rust
// Comprehensive tests for autodiff operations
|
|
// Following strict TDD - tests written BEFORE implementation fixes
|
|
|
|
use rtx_autograd::{Variable, Graph};
|
|
use rtx_tensor::{Tensor, Device};
|
|
use std::sync::Arc;
|
|
|
|
#[cfg(test)]
|
|
mod autodiff_gradient_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_backward_grad_scalar() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Create a scalar variable
|
|
let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
|
|
let y = &x * &x; // y = x^2
|
|
|
|
// Compute gradient dy/dx = 2x = 4
|
|
let grad = y.backward_grad(&x);
|
|
assert_eq!(grad.shape(), &[]);
|
|
|
|
let grad_value: f32 = grad.item();
|
|
assert!((grad_value - 4.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_backward_grad_vector() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Create vector variable
|
|
let x = Variable::new(Tensor::ones(&[3], &device), true);
|
|
let y = &x * &x + &x; // y = x^2 + x
|
|
|
|
// Compute gradient dy/dx = 2x + 1
|
|
let grad = y.backward_grad(&x);
|
|
assert_eq!(grad.shape(), &[3]);
|
|
|
|
let grad_data: Vec<f32> = grad.to_vec();
|
|
for val in grad_data {
|
|
assert!((val - 3.0).abs() < 1e-6); // 2*1 + 1 = 3
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_backward_grad_matrix() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Create matrix variable
|
|
let x = Variable::new(Tensor::ones(&[2, 3], &device), true);
|
|
let w = Variable::new(Tensor::full(&[3, 2], 2.0, &device), false);
|
|
|
|
// y = x @ w (matrix multiplication)
|
|
let y = x.matmul(&w);
|
|
|
|
// Compute gradient dy/dx
|
|
let grad_x = y.backward_grad(&x);
|
|
assert_eq!(grad_x.shape(), &[2, 3]);
|
|
|
|
// Each element should be 4.0 (sum of row in w = 2+2)
|
|
let grad_data: Vec<f32> = grad_x.to_vec();
|
|
for val in grad_data {
|
|
assert!((val - 4.0).abs() < 1e-6);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_backward_grad_chain_rule() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Test chain rule: z = sin(x^2)
|
|
let x = Variable::new(Tensor::scalar(std::f32::consts::PI / 4.0, DType::F32, &device), true);
|
|
let x_squared = &x * &x;
|
|
let z = x_squared.sin();
|
|
|
|
// dz/dx = cos(x^2) * 2x
|
|
let grad = z.backward_grad(&x);
|
|
|
|
let x_val = std::f32::consts::PI / 4.0;
|
|
let expected = (x_val * x_val).cos() * 2.0 * x_val;
|
|
|
|
let grad_value: f32 = grad.item();
|
|
assert!((grad_value - expected).abs() < 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_mean_square_operation() {
|
|
let device = Device::cpu();
|
|
|
|
// Test mean square operation
|
|
let data = vec![1.0, 2.0, 3.0, 4.0];
|
|
let x = Variable::new(Tensor::from_vec(data.clone(), &[4], &device), true);
|
|
|
|
// mean_square should compute mean(x^2)
|
|
let result = x.mean_square();
|
|
|
|
// Expected: (1 + 4 + 9 + 16) / 4 = 30 / 4 = 7.5
|
|
let value: f32 = result.value().item();
|
|
assert!((value - 7.5).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_value_extraction() {
|
|
let device = Device::cpu();
|
|
|
|
let data = vec![1.0, 2.0, 3.0];
|
|
let tensor = Tensor::from_vec(data.clone(), &[3], &device);
|
|
let var = Variable::new(tensor, false);
|
|
|
|
// Test value extraction
|
|
let tensor_value = var.value();
|
|
assert_eq!(tensor_value.shape(), &[3]);
|
|
|
|
let extracted: Vec<f32> = tensor_value.to_vec();
|
|
assert_eq!(extracted, data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradient_accumulation() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Test gradient accumulation through multiple paths
|
|
let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
|
|
|
|
// Two paths: y1 = x^2, y2 = 3x
|
|
let y1 = &x * &x;
|
|
let y2 = &x * &Variable::new(Tensor::scalar(3.0, DType::F32, &device), false);
|
|
|
|
// z = y1 + y2 = x^2 + 3x
|
|
let z = &y1 + &y2;
|
|
|
|
// dz/dx = 2x + 3 = 7
|
|
let grad = z.backward_grad(&x);
|
|
let grad_value: f32 = grad.item();
|
|
assert!((grad_value - 7.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stop_gradient() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Test gradient flow control
|
|
let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
|
|
let y = Variable::new(Tensor::scalar(3.0, DType::F32, &device), false); // requires_grad = false
|
|
|
|
let z = &x * &y;
|
|
|
|
// Should compute gradient w.r.t x
|
|
let grad_x = z.backward_grad(&x);
|
|
let grad_x_value: f32 = grad_x.item();
|
|
assert!((grad_x_value - 3.0).abs() < 1e-6);
|
|
|
|
// Should not compute gradient w.r.t y (requires_grad = false)
|
|
// This should return zeros or handle gracefully
|
|
let grad_y = z.backward_grad(&y);
|
|
let grad_y_value: f32 = grad_y.item();
|
|
assert!((grad_y_value - 0.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_complex_computational_graph() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Build a more complex graph: f(x,y) = (x^2 + y) * sin(x*y)
|
|
let x = Variable::new(Tensor::scalar(1.0, DType::F32, &device), true);
|
|
let y = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
|
|
|
|
let x_squared = &x * &x;
|
|
let sum_term = &x_squared + &y;
|
|
let product = &x * &y;
|
|
let sin_term = product.sin();
|
|
let result = &sum_term * &sin_term;
|
|
|
|
// Compute gradients
|
|
let grad_x = result.backward_grad(&x);
|
|
let grad_y = result.backward_grad(&y);
|
|
|
|
// Verify gradients are computed
|
|
assert_eq!(grad_x.shape(), &[]);
|
|
assert_eq!(grad_y.shape(), &[]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_batch_gradient_computation() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Test gradient computation for batch operations
|
|
let batch_size = 4;
|
|
let features = 3;
|
|
|
|
let x = Variable::new(Tensor::randn(&[batch_size, features], &device), true);
|
|
let w = Variable::new(Tensor::randn(&[features, 2], &device), true);
|
|
let b = Variable::new(Tensor::randn(&[2], &device), true);
|
|
|
|
// Linear layer: y = xW + b
|
|
let y = x.matmul(&w).add(&b);
|
|
|
|
// Compute mean loss
|
|
let loss = y.mean();
|
|
|
|
// Compute gradients
|
|
let grad_w = loss.backward_grad(&w);
|
|
let grad_b = loss.backward_grad(&b);
|
|
|
|
assert_eq!(grad_w.shape(), &[features, 2]);
|
|
assert_eq!(grad_b.shape(), &[2]);
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod autodiff_optimization_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_gradient_descent_step() {
|
|
let device = Device::cpu();
|
|
|
|
// Simple quadratic: f(x) = (x - 3)^2
|
|
let mut x = Variable::new(Tensor::scalar(0.0, DType::F32, &device), true);
|
|
let target = Variable::new(Tensor::scalar(3.0, DType::F32, &device), false);
|
|
|
|
let learning_rate = 0.1;
|
|
|
|
for _ in 0..20 {
|
|
// Compute loss
|
|
let diff = &x - ⌖
|
|
let loss = &diff * &diff;
|
|
|
|
// Compute gradient
|
|
let grad = loss.backward_grad(&x);
|
|
|
|
// Update parameter
|
|
let x_tensor = x.value();
|
|
let grad_tensor = grad;
|
|
let updated = x_tensor - &(grad_tensor * learning_rate);
|
|
x = Variable::new(updated, true);
|
|
}
|
|
|
|
// Should converge close to 3.0
|
|
let final_value: f32 = x.value().item();
|
|
assert!((final_value - 3.0).abs() < 0.1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multi_variable_optimization() {
|
|
let device = Device::cpu();
|
|
|
|
// Minimize f(x,y) = x^2 + y^2 + 2x - 4y + 5
|
|
// Minimum at x = -1, y = 2
|
|
let mut x = Variable::new(Tensor::scalar(5.0, DType::F32, &device), true);
|
|
let mut y = Variable::new(Tensor::scalar(-3.0, DType::F32, &device), true);
|
|
|
|
let learning_rate = 0.1;
|
|
|
|
for _ in 0..50 {
|
|
// Compute loss
|
|
let x_squared = &x * &x;
|
|
let y_squared = &y * &y;
|
|
let linear_x = &x * &Variable::new(Tensor::scalar(2.0, DType::F32, &device), false);
|
|
let linear_y = &y * &Variable::new(Tensor::scalar(-4.0, DType::F32, &device), false);
|
|
let constant = Variable::new(Tensor::scalar(5.0, DType::F32, &device), false);
|
|
|
|
let loss = &(&(&x_squared + &y_squared) + &linear_x) + &(&linear_y + &constant);
|
|
|
|
// Compute gradients
|
|
let grad_x = loss.backward_grad(&x);
|
|
let grad_y = loss.backward_grad(&y);
|
|
|
|
// Update parameters
|
|
let x_updated = x.value() - &(grad_x * learning_rate);
|
|
let y_updated = y.value() - &(grad_y * learning_rate);
|
|
|
|
x = Variable::new(x_updated, true);
|
|
y = Variable::new(y_updated, true);
|
|
}
|
|
|
|
// Should converge close to (-1, 2)
|
|
let x_final: f32 = x.value().item();
|
|
let y_final: f32 = y.value().item();
|
|
|
|
assert!((x_final - (-1.0)).abs() < 0.1);
|
|
assert!((y_final - 2.0).abs() < 0.1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_higher_order_derivatives() {
|
|
let graph = Graph::new();
|
|
let device = Device::cpu();
|
|
|
|
// Test second-order derivatives
|
|
let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
|
|
|
|
// f(x) = x^3
|
|
let x_squared = &x * &x;
|
|
let x_cubed = &x_squared * &x;
|
|
|
|
// First derivative: df/dx = 3x^2 = 12
|
|
let grad1 = x_cubed.backward_grad(&x);
|
|
let grad1_value: f32 = grad1.item();
|
|
assert!((grad1_value - 12.0).abs() < 1e-5);
|
|
|
|
// For second derivative, we'd need to differentiate grad1
|
|
// This tests that the gradient computation infrastructure works
|
|
}
|
|
|
|
#[test]
|
|
fn test_numerical_stability() {
|
|
let device = Device::cpu();
|
|
|
|
// Test with very small and very large values
|
|
let small = Variable::new(Tensor::scalar(1e-8, DType::F32, &device), true);
|
|
let large = Variable::new(Tensor::scalar(1e8, DType::F32, &device), true);
|
|
|
|
// Operations should not overflow or underflow
|
|
let result_small = &small * &small;
|
|
let result_large = &large + &large;
|
|
|
|
// Gradients should be computable
|
|
let grad_small = result_small.backward_grad(&small);
|
|
let grad_large = result_large.backward_grad(&large);
|
|
|
|
// Verify no NaN or Inf
|
|
let grad_small_val: f32 = grad_small.item();
|
|
let grad_large_val: f32 = grad_large.item();
|
|
|
|
assert!(grad_small_val.is_finite());
|
|
assert!(grad_large_val.is_finite());
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod integration_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_neural_network_forward_backward() {
|
|
let device = Device::cpu();
|
|
|
|
// Simple 2-layer network
|
|
let batch_size = 2;
|
|
let input_dim = 3;
|
|
let hidden_dim = 4;
|
|
let output_dim = 1;
|
|
|
|
// Initialize weights and biases
|
|
let w1 = Variable::new(Tensor::randn(&[input_dim, hidden_dim], &device) * 0.1, true);
|
|
let b1 = Variable::new(Tensor::zeros(&[hidden_dim], &device), true);
|
|
let w2 = Variable::new(Tensor::randn(&[hidden_dim, output_dim], &device) * 0.1, true);
|
|
let b2 = Variable::new(Tensor::zeros(&[output_dim], &device), true);
|
|
|
|
// Input and target
|
|
let x = Variable::new(Tensor::randn(&[batch_size, input_dim], &device), false);
|
|
let target = Variable::new(Tensor::ones(&[batch_size, output_dim], &device), false);
|
|
|
|
// Forward pass
|
|
let h1 = x.matmul(&w1).add(&b1).relu();
|
|
let output = h1.matmul(&w2).add(&b2);
|
|
|
|
// Loss (MSE)
|
|
let diff = &output - ⌖
|
|
let squared = &diff * &diff;
|
|
let loss = squared.mean();
|
|
|
|
// Backward pass
|
|
let grad_w1 = loss.backward_grad(&w1);
|
|
let grad_b1 = loss.backward_grad(&b1);
|
|
let grad_w2 = loss.backward_grad(&w2);
|
|
let grad_b2 = loss.backward_grad(&b2);
|
|
|
|
// Verify gradient shapes
|
|
assert_eq!(grad_w1.shape(), &[input_dim, hidden_dim]);
|
|
assert_eq!(grad_b1.shape(), &[hidden_dim]);
|
|
assert_eq!(grad_w2.shape(), &[hidden_dim, output_dim]);
|
|
assert_eq!(grad_b2.shape(), &[output_dim]);
|
|
}
|
|
} |