mirror of
https://github.com/huggingface/candle.git
synced 2025-06-16 10:38:54 +00:00
Use the same default as pytorch for sum. (#164)
This commit is contained in:
@ -19,7 +19,7 @@ fn simple_grad(device: &Device) -> Result<()> {
|
||||
fn sum_grad(device: &Device) -> Result<()> {
|
||||
let x = Var::new(&[3f32, 1., 4.], device)?;
|
||||
let x = x.as_tensor();
|
||||
let y = (x.sqr()?.sum(&[0])? * 2.)?;
|
||||
let y = (x.sqr()?.sum_keepdim(&[0])? * 2.)?;
|
||||
let grads = y.backward()?;
|
||||
let grad_x = grads.get(x).context("no grad for x")?;
|
||||
assert_eq!(y.to_vec1::<f32>()?, [52.]);
|
||||
@ -27,7 +27,7 @@ fn sum_grad(device: &Device) -> Result<()> {
|
||||
assert_eq!(grad_x.to_vec1::<f32>()?, &[12., 4., 16.]);
|
||||
|
||||
// Same test as before but squeezing on the last dimension.
|
||||
let y = (x.sqr()?.sum(&[0])? * 2.)?.squeeze(0)?;
|
||||
let y = (x.sqr()?.sum_keepdim(&[0])? * 2.)?.squeeze(0)?;
|
||||
let grads = y.backward()?;
|
||||
let grad_x = grads.get(x).context("no grad for x")?;
|
||||
assert_eq!(y.to_scalar::<f32>()?, 52.);
|
||||
|
Reference in New Issue
Block a user