Files
foxhunt/crates/ml/tests/test_var_gradient_flow.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

76 lines
2.7 KiB
Rust

//! Minimal test to understand Var vs Tensor gradient flow with scatter_add
use candle_core::{Device, DType, Tensor, Var};
use ml::MLError;
#[test]
fn test_scatter_add_tensor_vs_var() -> Result<(), MLError> {
let device = Device::cuda_if_available(0)?;
println!("\n=== Testing Tensor::zeros + scatter_add ===");
{
// Input that should have gradients
let input = Tensor::ones((2, 3), DType::F32, &device)?;
// Base for scatter (using Tensor::zeros)
let base = Tensor::zeros((2, 3), DType::F32, &device)?;
let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;
// Scatter
let result = base.scatter_add(&indices, &input, 1)?;
// Compute loss and backward
let loss = result.sum_all()?;
let grads = loss.backward()?;
let has_grads = grads.get(&input).is_some();
println!("Tensor::zeros -> has gradients: {}", has_grads);
}
println!("\n=== Testing Var::zeros + scatter_add ===");
{
// Input that should have gradients
let input = Tensor::ones((2, 3), DType::F32, &device)?;
// Base for scatter (using Var::zeros)
let base_var = Var::zeros((2, 3), DType::F32, &device)?;
let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;
// Scatter - need to convert Var to Tensor for scatter_add
let base_tensor = base_var.as_tensor();
let result = base_tensor.scatter_add(&indices, &input, 1)?;
// Compute loss and backward
let loss = result.sum_all()?;
let grads = loss.backward()?;
let has_grads = grads.get(&input).is_some();
println!("Var::zeros -> has gradients: {}", has_grads);
}
println!("\n=== Testing Var::from_tensor (input) + scatter_add ===");
{
// Input wrapped in Var
let input_tensor = Tensor::ones((2, 3), DType::F32, &device)?;
let input_var = Var::from_tensor(&input_tensor)?;
// Base for scatter (regular Tensor)
let base = Tensor::zeros((2, 3), DType::F32, &device)?;
let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;
// Scatter - use Var as Tensor
let result = base.scatter_add(&indices, &input_var.as_tensor(), 1)?;
// Compute loss and backward
let loss = result.sum_all()?;
let grads = loss.backward()?;
let has_grads_tensor = grads.get(&input_tensor).is_some();
let has_grads_var = grads.get(input_var.as_tensor()).is_some();
println!("Var(input) -> has gradients on tensor: {}", has_grads_tensor);
println!("Var(input) -> has gradients on var: {}", has_grads_var);
}
Ok(())
}