Files
foxhunt/crates/ml/tests/test_var_source_gradients.rs
jgrusewski 9c3d741a08 refactor: restructure repo — crates/, bin/, testing/ layout
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>
2026-02-25 11:56:00 +01:00

63 lines
2.3 KiB
Rust

//! Test gradient flow when source is derived from a Var
use candle_core::{Device, DType, Tensor, Var};
use ml::MLError;
#[test]
fn test_simple_var_gradient_flow() -> Result<(), MLError> {
let device = Device::cuda_if_available(0)?;
println!("\n=== Test 1: Simple Var multiplication ===");
{
let x_var = Var::new(&[1.0f32, 2.0, 3.0], &device)?;
let y = (x_var.as_tensor() * 2.0)?;
let loss = y.sum_all()?;
let grads = loss.backward()?;
let has_grads = grads.get(x_var.as_tensor()).is_some();
println!("Var multiplication: has gradients: {}", has_grads);
}
println!("\n=== Test 2: Var used in scatter_add source ===");
{
let source_var = Var::new(&[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], &device)?;
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_var.as_tensor(), 1)?;
let loss = result.sum_all()?;
let grads = loss.backward()?;
let has_grads = grads.get(source_var.as_tensor()).is_some();
println!("scatter_add with Var source: has gradients: {}", has_grads);
if let Some(grad) = grads.get(source_var.as_tensor()) {
let grad_sum: f32 = grad.sum_all()?.to_scalar()?;
println!(" Gradient sum: {}", grad_sum);
}
}
println!("\n=== Test 3: Derived tensor from Var in scatter_add ===");
{
let base_var = Var::new(&[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], &device)?;
let source = (base_var.as_tensor() * 2.0)?; // Derive from Var
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(base_var.as_tensor()).is_some();
println!("scatter_add with derived source: has gradients on base_var: {}", has_grads);
if let Some(grad) = grads.get(base_var.as_tensor()) {
let grad_sum: f32 = grad.sum_all()?.to_scalar()?;
println!(" Gradient sum: {}", grad_sum);
}
}
Ok(())
}