huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 10:26:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

More QOL changes, binary op for constants.

Browse Source
This commit is contained in:
laurent
2023-06-21 08:59:08 +01:00
parent 0839954770
commit f319583530
2 changed files with 22 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/tensor.rs
View File

@ -445,7 +445,7 @@ impl Tensor {
}
macro_rules! bin_trait {
($trait:ident, $fn1:ident) => {
($trait:ident, $fn1:ident, $mul:expr, $add:expr) => {
impl<B: std::borrow::Borrow<Tensor>> std::ops::$trait<B> for Tensor {
type Output = Result<Tensor>;
@ -477,10 +477,26 @@ macro_rules! bin_trait {
Tensor::$fn1(&self, rhs?.borrow())
}
}
impl std::ops::$trait<f64> for Tensor {
type Output = Result<Tensor>;
fn $fn1(self, rhs: f64) -> Self::Output {
self.affine($mul(rhs), $add(rhs))
}
}
impl std::ops::$trait<f64> for &Tensor {
type Output = Result<Tensor>;
fn $fn1(self, rhs: f64) -> Self::Output {
self.affine($mul(rhs), $add(rhs))
}
}
};
}
bin_trait!(Add, add);
bin_trait!(Sub, sub);
bin_trait!(Mul, mul);
bin_trait!(Div, div);
bin_trait!(Add, add, |_| 1., |v| v);
bin_trait!(Sub, sub, |_| 1., |v: f64| -v);
bin_trait!(Mul, mul, |v| v, |_| 0.);
bin_trait!(Div, div, |v| 1. / v, |_| 0.);

2
tests/grad_tests.rs
View File

@ -4,7 +4,7 @@ use candle::{Device, Tensor};
#[test]
fn simple_grad() -> Result<()> {
let x = Tensor::var(&[3f32, 1., 4.], Device::Cpu)?;
let y = x.mul(&x)?.add(&x.affine(5., 0.)?)?.affine(1., 4.)?;
let y = (((&x * &x)? + &x * 5f64)? + 4f64)?;
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.]);