dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor: neg-ops (#70)

This commit is contained in:
Nathaniel Simard 2022-11-05 17:19:04 -04:00 committed by GitHub
parent 94b0283bac
commit ad23898d23
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
15 changed files with 86 additions and 209 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
burn-tensor/src/tensor/backend/autodiff/ops/erf.rs
View File

@ -12,7 +12,7 @@ register_ops!(
name ADTensorErfOps,
partial |state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>|{
let value = state.input.value();
let exponent = value.powf(2.0.to_elem()).neg();
let exponent = B::neg(&value.powf(2.0.to_elem()));
let numerator = B::mul_scalar(&exponent.exp(), &2.0.to_elem());
let denominator = std::f64::consts::PI.sqrt().to_elem();
let value = B::div_scalar(&numerator, &denominator);

1
burn-tensor/src/tensor/backend/autodiff/ops/mod.rs
View File

@ -11,7 +11,6 @@ mod log;
mod map_comparison;
mod mask;
mod module;
mod neg;
mod pow;
mod precision;
mod reshape;

49
burn-tensor/src/tensor/backend/autodiff/ops/neg.rs
View File

@ -1,49 +0,0 @@
use crate::tensor::backend::Backend;
use crate::{
execute_ops,
graph::ops::{UnaryOps, UnaryOpsNodeState},
register_ops,
tensor::{backend::autodiff::ADTensor, ops::*},
};
register_ops!(
ops UnaryOps,
name ADTensorNegOps,
partial |state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>|{
state.output.grad().neg()
},
);
impl<B: Backend, P, const D: usize> TensorOpsNeg<P, D> for ADTensor<D, B> {
fn neg(&self) -> Self {
execute_ops!(
input self.node.clone(),
out TensorOpsNeg::neg(&self.tensor()),
ops ADTensorNegOps::<B, D>::new(),
)
}
}
#[cfg(test)]
mod tests {
use crate::tensor::{backend::autodiff::helper::TestADTensor, Data};
#[test]
fn should_diff_neg() {
let data_1 = Data::<f64, 2>::from([[1.0, 7.0], [2.0, 3.0]]);
let data_2 = Data::<f64, 2>::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2.neg());
let tensor_4 = tensor_3.neg();
let grads = tensor_4.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(grad_1.to_data(), Data::from([[11.0, 5.0], [11.0, 5.0]]));
assert_eq!(grad_2.to_data(), Data::from([[3.0, 3.0], [10.0, 10.0]]));
}
}

31
burn-tensor/src/tensor/backend/autodiff/ops/tensor.rs
View File

@ -5,7 +5,7 @@ use crate::{
Backend,
},
graph::ops::{BinaryOps, BinaryOpsNodeState, UnaryOps, UnaryOpsNodeState},
ops::{Ones, TensorOps, TensorOpsNeg, TensorOpsTranspose},
ops::{Ones, TensorOps, TensorOpsTranspose},
Data, Shape,
};
@ -195,7 +195,7 @@ impl<B: Backend> TensorOps<ADBackendDecorator<B>> for ADBackendDecorator<B> {
B::TensorPrimitive<D>,
>,
) -> B::TensorPrimitive<D> {
state.output.grad().neg()
B::neg(&state.output.grad())
}
}
@ -336,7 +336,7 @@ impl<B: Backend> TensorOps<ADBackendDecorator<B>> for ADBackendDecorator<B> {
) -> B::TensorPrimitive<D> {
let value_left = state.left.value();
let value_right = state.right.value();
let value = B::div(&value_left.neg(), &B::mul(&value_right, &value_right));
let value = B::div(&B::neg(&value_left), &B::mul(&value_right, &value_right));
B::mul(&state.output.grad(), &value)
}
@ -422,4 +422,29 @@ impl<B: Backend> TensorOps<ADBackendDecorator<B>> for ADBackendDecorator<B> {
binary_ops_wrapper(lhs.node.clone(), rhs.node.clone(), output, ops)
}
fn neg<const D: usize>(
tensor: &<ADBackendDecorator<B> as Backend>::TensorPrimitive<D>,
) -> <ADBackendDecorator<B> as Backend>::TensorPrimitive<D> {
#[derive(Default, Debug)]
struct NegBackward<B: Backend, const D: usize> {
_b: B,
}
impl<B: Backend, const D: usize> UnaryOps<B::TensorPrimitive<D>, B::TensorPrimitive<D>>
for NegBackward<B, D>
{
fn partial(
&self,
state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>,
) -> B::TensorPrimitive<D> {
B::neg(&state.output.grad())
}
}
let output = B::neg(tensor.tensor_ref());
let ops = NegBackward::<B, D>::default();
unary_ops_wrapper(tensor.node.clone(), output, ops)
}
}

1
burn-tensor/src/tensor/backend/base.rs
View File

@ -23,7 +23,6 @@ pub trait Backend:
type IntegerBackend: Backend<Elem = i64, Device = Self::Device>;
type TensorPrimitive<const D: usize>: std::ops::Add<Self::TensorPrimitive<D>, Output = Self::TensorPrimitive<D>>
+ TensorOpsTranspose<Self::Elem, D>
+ TensorOpsNeg<Self::Elem, D>
+ TensorOpsDetach<Self::Elem, D>
+ Zeros<Self::TensorPrimitive<D>>
+ Ones<Self::TensorPrimitive<D>>

1
burn-tensor/src/tensor/backend/ndarray/ops/mod.rs
View File

@ -9,7 +9,6 @@ mod index;
mod log;
mod map_comparison;
mod mask;
mod neg;
mod pow;
mod precision;
mod reshape;

27
burn-tensor/src/tensor/backend/ndarray/ops/neg.rs
View File

@ -1,27 +0,0 @@
use crate::tensor::{backend::ndarray::NdArrayTensor, ops::*};
use ndarray::{LinalgScalar, ScalarOperand};
impl<P, const D: usize> TensorOpsNeg<P, D> for NdArrayTensor<P, D>
where
P: Clone + LinalgScalar + Default + std::fmt::Debug + ScalarOperand,
{
fn neg(&self) -> Self {
let minus_one = P::zero() - P::one();
let array = self.array.clone() * minus_one;
let array = array.into_shared();
let shape = self.shape;
Self { array, shape }
}
}
impl<P, const D: usize> std::ops::Neg for NdArrayTensor<P, D>
where
P: Clone + LinalgScalar + Default + std::fmt::Debug + ScalarOperand,
{
type Output = Self;
fn neg(self) -> Self::Output {
TensorOpsNeg::neg(&self)
}
}

8
burn-tensor/src/tensor/backend/ndarray/tensor_ops.rs
View File

@ -2,7 +2,7 @@ use super::{BatchMatrix, NdArrayBackend, NdArrayTensor};
use crate::{
backend::{Backend, NdArrayDevice},
ops::TensorOps,
Data, NdArrayElement, Shape,
Data, ElementConversion, NdArrayElement, Shape,
};
impl<E: NdArrayElement> TensorOps<NdArrayBackend<E>> for NdArrayBackend<E> {
@ -172,4 +172,10 @@ impl<E: NdArrayElement> TensorOps<NdArrayBackend<E>> for NdArrayBackend<E> {
NdArrayTensor::from_bmatrix(output)
}
fn neg<const D: usize>(
tensor: &NdArrayTensor<E, D>,
) -> <NdArrayBackend<E> as Backend>::TensorPrimitive<D> {
Self::mul_scalar(tensor, &(-1f32).to_elem::<E>())
}
}

1
burn-tensor/src/tensor/backend/tch/ops/mod.rs
View File

@ -9,7 +9,6 @@ mod index;
mod log;
mod map_comparison;
mod mask;
mod neg;
mod pow;
mod precision;
mod reshape;

44
burn-tensor/src/tensor/backend/tch/ops/neg.rs
View File

@ -1,44 +0,0 @@
use crate::tensor::{backend::tch::TchTensor, ops::*};
impl<P: tch::kind::Element + Default + Copy + std::fmt::Debug, const D: usize> TensorOpsNeg<P, D>
for TchTensor<P, D>
{
fn neg(&self) -> Self {
let tensor = -(&self.tensor);
let kind = self.kind;
let shape = self.shape;
Self {
tensor,
shape,
kind,
}
}
}
impl<P: tch::kind::Element + Default + std::fmt::Debug + Copy, const D: usize> std::ops::Neg
for TchTensor<P, D>
{
type Output = Self;
fn neg(self) -> Self::Output {
TensorOpsNeg::neg(&self)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tensor::Data;
#[test]
fn should_support_neg_ops() {
let data = Data::<f64, 2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
let tensor = TchTensor::from_data(data, tch::Device::Cpu);
let data_actual = tensor.neg().into_data();
let data_expected = Data::from([[-0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]);
assert_eq!(data_expected, data_actual);
}
}

101
burn-tensor/src/tensor/backend/tch/tensor_ops.rs
View File

@ -1,7 +1,6 @@
use std::ops::{Add, Div, Mul, Sub};
use super::{TchBackend, TchDevice, TchKind, TchTensor};
use crate::{backend::Backend, ops::TensorOps, Data, Shape, TchElement};
use crate::{backend::Backend, ops::TensorOps, Data, ElementConversion, Shape, TchElement};
use std::ops::{Add, Div, Mul, Sub};
impl<E: TchElement> TensorOps<TchBackend<E>> for TchBackend<E> {
fn shape<const D: usize>(tensor: &<TchBackend<E> as Backend>::TensorPrimitive<D>) -> &Shape<D> {
@ -64,126 +63,78 @@ impl<E: TchElement> TensorOps<TchBackend<E>> for TchBackend<E> {
shape: Shape<D>,
device: <TchBackend<E> as Backend>::Device,
) -> <TchBackend<E> as Backend>::TensorPrimitive<D> {
let kind = TchKind::new();
let kind = TchKind::<E>::new();
let tensor =
tch::Tensor::empty(&shape.dims.map(|a| a as i64), (kind.kind(), device.into()));
TchTensor {
kind,
tensor,
shape,
}
to_tensor(tensor)
}
fn add<const D: usize>(lhs: &TchTensor<E, D>, rhs: &TchTensor<E, D>) -> TchTensor<E, D> {
let tensor = (&lhs.tensor).add(&rhs.tensor);
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn add_scalar<const D: usize>(lhs: &TchTensor<E, D>, rhs: &E) -> TchTensor<E, D> {
let other: f64 = (rhs.clone()).to_elem();
let tensor = (&lhs.tensor).add(other).to_kind(lhs.kind.kind());
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn sub<const D: usize>(lhs: &TchTensor<E, D>, rhs: &TchTensor<E, D>) -> TchTensor<E, D> {
let tensor = (&lhs.tensor).sub(&rhs.tensor);
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn sub_scalar<const D: usize>(lhs: &TchTensor<E, D>, rhs: &E) -> TchTensor<E, D> {
let other: f64 = (rhs.clone()).to_elem();
let tensor = (&lhs.tensor).sub(other).to_kind(lhs.kind.kind());
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn mul<const D: usize>(lhs: &TchTensor<E, D>, rhs: &TchTensor<E, D>) -> TchTensor<E, D> {
let tensor = (&lhs.tensor).mul(&rhs.tensor);
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn mul_scalar<const D: usize>(lhs: &TchTensor<E, D>, rhs: &E) -> TchTensor<E, D> {
let other: f64 = (rhs.clone()).to_elem();
let tensor = (&lhs.tensor).mul(other).to_kind(lhs.kind.kind());
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn div<const D: usize>(lhs: &TchTensor<E, D>, rhs: &TchTensor<E, D>) -> TchTensor<E, D> {
let tensor = (&lhs.tensor).div(&rhs.tensor);
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn div_scalar<const D: usize>(lhs: &TchTensor<E, D>, rhs: &E) -> TchTensor<E, D> {
let other: f64 = (rhs.clone()).to_elem();
let tensor = (&lhs.tensor).div(other).to_kind(lhs.kind.kind());
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind,
}
to_tensor(tensor)
}
fn matmul<const D: usize>(lhs: &TchTensor<E, D>, rhs: &TchTensor<E, D>) -> TchTensor<E, D> {
let tensor = lhs.tensor.matmul(&rhs.tensor);
let kind = lhs.kind;
let shape = Shape::from(tensor.size());
to_tensor(tensor)
}
TchTensor {
tensor,
shape,
kind,
}
fn neg<const D: usize>(tensor: &TchTensor<E, D>) -> TchTensor<E, D> {
Self::mul_scalar(tensor, &(-1f32).to_elem::<E>())
}
}
fn to_tensor<const D: usize, E: TchElement>(tensor: tch::Tensor) -> TchTensor<E, D> {
let shape = Shape::from(tensor.size());
TchTensor {
tensor,
shape,
kind: TchKind::new(),
}
}

2
burn-tensor/src/tensor/base.rs
View File

@ -224,7 +224,7 @@ where
///
/// `y = -x`
pub fn neg(&self) -> Self {
Self::new(self.value.neg())
Self::new(B::neg(&self.value))
}
/// Applies element wise multiplication operation.

5
burn-tensor/src/tensor/ops/base.rs
View File

@ -102,6 +102,7 @@ pub trait TensorOps<B: Backend> {
lhs: &B::TensorPrimitive<D>,
rhs: &B::TensorPrimitive<D>,
) -> B::TensorPrimitive<D>;
fn neg<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;
}
pub trait TensorOpsTranspose<E, const D: usize> {
@ -109,10 +110,6 @@ pub trait TensorOpsTranspose<E, const D: usize> {
fn swap_dims(&self, dim1: usize, dim2: usize) -> Self;
}
pub trait TensorOpsNeg<E, const D: usize> {
fn neg(&self) -> Self;
}
pub trait TensorOpsReshape<B: Backend, const D: usize> {
fn reshape<const D2: usize>(&self, shape: Shape<D2>) -> B::TensorPrimitive<D2>;
}

1
burn-tensor/tests/tensor/grad/mod.rs
View File

@ -4,5 +4,6 @@ mod cross_entropy;
mod div;
mod matmul;
mod mul;
mod neg;
mod softmax;
mod sub;

21
burn-tensor/tests/tensor/grad/neg.rs Normal file
View File

@ -0,0 +1,21 @@
use crate::tensor::TestADTensor;
use burn_tensor::Data;
#[test]
fn should_diff_neg() {
let data_1 = Data::<f32, 2>::from([[1.0, 7.0], [2.0, 3.0]]);
let data_2 = Data::<f32, 2>::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2.neg());
let tensor_4 = tensor_3.neg();
let grads = tensor_4.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(grad_1.to_data(), Data::from([[11.0, 5.0], [11.0, 5.0]]));
assert_eq!(grad_2.to_data(), Data::from([[3.0, 3.0], [10.0, 10.0]]));
}