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>
75 lines
2.8 KiB
Rust
75 lines
2.8 KiB
Rust
//! Test gradient flow through scatter_add SOURCE parameter
|
|
|
|
use candle_core::{Device, DType, Tensor, Var};
|
|
use ml::MLError;
|
|
|
|
#[test]
|
|
fn test_scatter_add_source_gradients() -> Result<(), MLError> {
|
|
let device = Device::cuda_if_available(0)?;
|
|
|
|
println!("\n=== Testing gradient flow through scatter_add SOURCE (values to scatter) ===");
|
|
|
|
// Create source values that should receive gradients
|
|
let source_base = Tensor::ones((2, 3), DType::F32, &device)?;
|
|
let source = (&source_base * 2.0)?; // source = source_base * 2
|
|
|
|
println!("Source tensor created (should track gradients back to source_base)");
|
|
|
|
// Base for scatter - TEST 1: using Tensor::zeros
|
|
println!("\n--- Test 1: Tensor::zeros as base ---");
|
|
{
|
|
let base = Tensor::zeros((2, 3), DType::F32, &device)?;
|
|
let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;
|
|
|
|
let result = base.scatter_add(&indices, &source, 1)?;
|
|
let loss = result.sum_all()?;
|
|
let grads = loss.backward()?;
|
|
|
|
let has_grads = grads.get(&source_base).is_some();
|
|
println!("Gradients flow to source_base: {}", has_grads);
|
|
if let Some(grad) = grads.get(&source_base) {
|
|
let grad_sum: f32 = grad.sum_all()?.to_scalar()?;
|
|
println!(" Gradient sum: {}", grad_sum);
|
|
}
|
|
}
|
|
|
|
// TEST 2: using Var::zeros as base
|
|
println!("\n--- Test 2: Var::zeros as base ---");
|
|
{
|
|
let base_var = Var::zeros((2, 3), DType::F32, &device)?;
|
|
let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;
|
|
|
|
let result = base_var.as_tensor().scatter_add(&indices, &source, 1)?;
|
|
let loss = result.sum_all()?;
|
|
let grads = loss.backward()?;
|
|
|
|
let has_grads = grads.get(&source_base).is_some();
|
|
println!("Gradients flow to source_base: {}", has_grads);
|
|
if let Some(grad) = grads.get(&source_base) {
|
|
let grad_sum: f32 = grad.sum_all()?.to_scalar()?;
|
|
println!(" Gradient sum: {}", grad_sum);
|
|
}
|
|
}
|
|
|
|
// TEST 3: using Var::from_tensor on zero tensor as base
|
|
println!("\n--- Test 3: Var::from_tensor(zeros) as base ---");
|
|
{
|
|
let zeros_tensor = Tensor::zeros((2, 3), DType::F32, &device)?;
|
|
let base_var = Var::from_tensor(&zeros_tensor)?;
|
|
let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;
|
|
|
|
let result = base_var.as_tensor().scatter_add(&indices, &source, 1)?;
|
|
let loss = result.sum_all()?;
|
|
let grads = loss.backward()?;
|
|
|
|
let has_grads = grads.get(&source_base).is_some();
|
|
println!("Gradients flow to source_base: {}", has_grads);
|
|
if let Some(grad) = grads.get(&source_base) {
|
|
let grad_sum: f32 = grad.sum_all()?.to_scalar()?;
|
|
println!(" Gradient sum: {}", grad_sum);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|