Simplify the test.

This commit is contained in:
laurent
2023-06-21 07:33:42 +01:00
parent 78bac0ed32
commit 23db8a7da8

View File

@ -4,13 +4,12 @@ use candle::{Device, Tensor};
#[test]
fn simple_grad() -> Result<()> {
let x = Tensor::var(&[3f32, 1., 4.], Device::Cpu)?;
let five = Tensor::new(&[5f32, 5., 5.], Device::Cpu)?;
let y = x.mul(&x)?.add(&x.mul(&five)?)?.add(&five)?;
let y = x.mul(&x)?.add(&x.affine(5., 0.)?)?.affine(1., 4.)?;
let grads = y.backward()?;
let grad_x = grads.get(&x.id()).context("no grad for x")?;
assert_eq!(x.to_vec1::<f32>()?, [3., 1., 4.]);
// y = x^2 + 5.x + 5
assert_eq!(y.to_vec1::<f32>()?, [29., 11., 41.]);
// y = x^2 + 5.x + 4
assert_eq!(y.to_vec1::<f32>()?, [28., 10., 40.]);
// dy/dx = 2.x + 5
assert_eq!(grad_x.to_vec1::<f32>()?, [11., 7., 13.]);
Ok(())