mirror of https://github.com/tracel-ai/burn.git
refactor: neg-ops (#70)
This commit is contained in:
parent
94b0283bac
commit
ad23898d23
|
@ -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);
|
||||
|
|
|
@ -11,7 +11,6 @@ mod log;
|
|||
mod map_comparison;
|
||||
mod mask;
|
||||
mod module;
|
||||
mod neg;
|
||||
mod pow;
|
||||
mod precision;
|
||||
mod reshape;
|
||||
|
|
|
@ -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]]));
|
||||
}
|
||||
}
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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>>
|
||||
|
|
|
@ -9,7 +9,6 @@ mod index;
|
|||
mod log;
|
||||
mod map_comparison;
|
||||
mod mask;
|
||||
mod neg;
|
||||
mod pow;
|
||||
mod precision;
|
||||
mod reshape;
|
||||
|
|
|
@ -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)
|
||||
}
|
||||
}
|
|
@ -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>())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -9,7 +9,6 @@ mod index;
|
|||
mod log;
|
||||
mod map_comparison;
|
||||
mod mask;
|
||||
mod neg;
|
||||
mod pow;
|
||||
mod precision;
|
||||
mod reshape;
|
||||
|
|
|
@ -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);
|
||||
}
|
||||
}
|
|
@ -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(),
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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.
|
||||
|
|
|
@ -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>;
|
||||
}
|
||||
|
|
|
@ -4,5 +4,6 @@ mod cross_entropy;
|
|||
mod div;
|
||||
mod matmul;
|
||||
mod mul;
|
||||
mod neg;
|
||||
mod softmax;
|
||||
mod sub;
|
||||
|
|
|
@ -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]]));
|
||||
}
|
Loading…
Reference in New Issue