refactor: sub ops (tracel-ai#66)

djdisodo · Nov 5, 2022 · 2bdad6f · 2bdad6f
1 parent 0f4c1e4
commit 2bdad6f
Show file tree

Hide file tree

Showing 17 changed files with 189 additions and 226 deletions.
diff --git a/README.md b/README.md
@@ -4,7 +4,7 @@
 [![Current Crates.io Version](https://img.shields.io/crates/v/burn.svg)](https://crates.io/crates/burn)
 [![Test Status](https://github.com/burn-rs/burn/actions/workflows/test-burn.yml/badge.svg)](https://github.com/burn-rs/burn/actions/workflows/test-burn.yml)
 [![Documentation](https://docs.rs/burn/badge.svg)](https://docs.rs/burn)
-[![Rust Version](https://img.shields.io/badge/Rust-1.65.0-blue)](https://releases.rs/docs/unreleased/1.65.0)
+[![Rust Version](https://img.shields.io/badge/Rust-1.65.0-blue)](https://releases.rs/docs/released/1.65.0)
 [![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/burn-rs/burn/blob/master/LICENSE)
 
 > This library aims to be a complete deep learning framework with extreme flexibility written in Rust. 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/mod.rs b/burn-tensor/src/tensor/backend/autodiff/ops/mod.rs
@@ -18,7 +18,6 @@ mod neg;
 mod pow;
 mod precision;
 mod reshape;
-mod sub;
 mod tensor;
 mod transpose;
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/sub.rs b/burn-tensor/src/tensor/backend/autodiff/ops/sub.rs
diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/tensor.rs b/burn-tensor/src/tensor/backend/autodiff/ops/tensor.rs
@@ -5,7 +5,7 @@ use crate::{
         Backend,
     },
     graph::ops::{BinaryOps, BinaryOpsNodeState, UnaryOps, UnaryOpsNodeState},
-    ops::TensorOps,
+    ops::{TensorOps, TensorOpsNeg},
     Data, Shape,
 };
 
@@ -65,6 +65,38 @@ impl<B: Backend, const D: usize>
     }
 }
 
+#[derive(Default, Debug)]
+struct SubBackward<B: Backend, const D: usize> {
+    _b: B,
+}
+
+impl<B: Backend, const D: usize>
+    BinaryOps<B::TensorPrimitive<D>, B::TensorPrimitive<D>, B::TensorPrimitive<D>>
+    for SubBackward<B, D>
+{
+    fn partial_left(
+        &self,
+        state: &BinaryOpsNodeState<
+            B::TensorPrimitive<D>,
+            B::TensorPrimitive<D>,
+            B::TensorPrimitive<D>,
+        >,
+    ) -> B::TensorPrimitive<D> {
+        state.output.grad()
+    }
+
+    fn partial_right(
+        &self,
+        state: &BinaryOpsNodeState<
+            B::TensorPrimitive<D>,
+            B::TensorPrimitive<D>,
+            B::TensorPrimitive<D>,
+        >,
+    ) -> B::TensorPrimitive<D> {
+        state.output.grad().neg()
+    }
+}
+
 #[derive(Default, Debug)]
 struct AddScalarBackward<B: Backend, const D: usize> {
     _b: B,
@@ -81,6 +113,22 @@ impl<B: Backend, const D: usize> UnaryOps<B::TensorPrimitive<D>, B::TensorPrimit
     }
 }
 
+#[derive(Default, Debug)]
+struct SubScalarBackward<B: Backend, const D: usize> {
+    _b: B,
+}
+
+impl<B: Backend, const D: usize> UnaryOps<B::TensorPrimitive<D>, B::TensorPrimitive<D>>
+    for SubScalarBackward<B, D>
+{
+    fn partial(
+        &self,
+        state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>,
+    ) -> B::TensorPrimitive<D> {
+        state.output.grad()
+    }
+}
+
 impl<B: Backend> TensorOps<ADBackendDecorator<B>> for ADBackendDecorator<B> {
     fn shape<const D: usize>(
         tensor: &<ADBackendDecorator<B> as Backend>::TensorPrimitive<D>,
@@ -162,4 +210,24 @@ impl<B: Backend> TensorOps<ADBackendDecorator<B>> for ADBackendDecorator<B> {
 
         unary_ops_wrapper(lhs.node.clone(), output, ops)
     }
+
+    fn sub<const D: usize>(
+        lhs: &<ADBackendDecorator<B> as Backend>::TensorPrimitive<D>,
+        rhs: &<ADBackendDecorator<B> as Backend>::TensorPrimitive<D>,
+    ) -> <ADBackendDecorator<B> as Backend>::TensorPrimitive<D> {
+        let output = B::sub(lhs.tensor_ref(), rhs.tensor_ref());
+        let ops = SubBackward::<B, D>::default();
+
+        binary_ops_wrapper(lhs.node.clone(), rhs.node.clone(), output, ops)
+    }
+
+    fn sub_scalar<const D: usize>(
+        lhs: &<ADBackendDecorator<B> as Backend>::TensorPrimitive<D>,
+        rhs: &<ADBackendDecorator<B> as Backend>::Elem,
+    ) -> <ADBackendDecorator<B> as Backend>::TensorPrimitive<D> {
+        let output = B::sub_scalar(lhs.tensor_ref(), rhs);
+        let ops = SubScalarBackward::<B, D>::default();
+
+        unary_ops_wrapper(lhs.node.clone(), output, ops)
+    }
 }
diff --git a/burn-tensor/src/tensor/backend/base.rs b/burn-tensor/src/tensor/backend/base.rs
@@ -27,7 +27,6 @@ pub trait Backend:
         + TensorOpsMul<Self::Elem, D>
         + TensorOpsDiv<Self::Elem, D>
         + TensorOpsNeg<Self::Elem, D>
-        + TensorOpsSub<Self::Elem, D>
         + TensorOpsDetach<Self::Elem, D>
         + Zeros<Self::TensorPrimitive<D>>
         + Ones<Self::TensorPrimitive<D>>

diff --git a/burn-tensor/src/tensor/backend/ndarray/ops/map_comparison.rs b/burn-tensor/src/tensor/backend/ndarray/ops/map_comparison.rs
@@ -8,7 +8,7 @@ where
     E: NdArrayElement,
 {
     fn equal(&self, other: &Self) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
-        let tensor = self.sub(other);
+        let tensor = NdArrayBackend::<E>::sub(self, other);
         let zero = E::zeros(&E::default());
         tensor.equal_scalar(&zero)
     }
@@ -26,7 +26,7 @@ where
     }
 
     fn greater(&self, other: &Self) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
-        let tensor = self.sub(other);
+        let tensor = NdArrayBackend::<E>::sub(self, other);
         let zero = E::zeros(&E::default());
         tensor.greater_scalar(&zero)
     }
@@ -47,7 +47,7 @@ where
         &self,
         other: &Self,
     ) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
-        let tensor = self.sub(other);
+        let tensor = NdArrayBackend::<E>::sub(self, other);
         let zero = E::zeros(&E::default());
         tensor.greater_equal_scalar(&zero)
     }
@@ -65,7 +65,7 @@ where
     }
 
     fn lower(&self, other: &Self) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
-        let tensor = self.sub(other);
+        let tensor = NdArrayBackend::<E>::sub(self, other);
         let zero = E::zeros(&E::default());
         tensor.lower_scalar(&zero)
     }
@@ -83,7 +83,7 @@ where
     }
 
     fn lower_equal(&self, other: &Self) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
-        let tensor = self.sub(other);
+        let tensor = NdArrayBackend::<E>::sub(self, other);
         let zero = E::zeros(&E::default());
         tensor.lower_equal_scalar(&zero)
     }
@@ -100,11 +100,3 @@ where
         }
     }
 }
-
-#[cfg(tests)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_greater() {}
-}
diff --git a/burn-tensor/src/tensor/backend/ndarray/ops/mod.rs b/burn-tensor/src/tensor/backend/ndarray/ops/mod.rs
@@ -16,5 +16,4 @@ mod neg;
 mod pow;
 mod precision;
 mod reshape;
-mod sub;
 mod transpose;
diff --git a/burn-tensor/src/tensor/backend/ndarray/ops/sub.rs b/burn-tensor/src/tensor/backend/ndarray/ops/sub.rs
diff --git a/burn-tensor/src/tensor/backend/ndarray/tensor_ops.rs b/burn-tensor/src/tensor/backend/ndarray/tensor_ops.rs
@@ -80,6 +80,25 @@ impl<E: NdArrayElement> TensorOps<NdArrayBackend<E>> for NdArrayBackend<E> {
         let array = lhs.array.clone() + *rhs;
         let shape = lhs.shape;
 
+        NdArrayTensor { array, shape }
+    }
+    fn sub<const D: usize>(
+        lhs: &<NdArrayBackend<E> as Backend>::TensorPrimitive<D>,
+        rhs: &<NdArrayBackend<E> as Backend>::TensorPrimitive<D>,
+    ) -> <NdArrayBackend<E> as Backend>::TensorPrimitive<D> {
+        let array = lhs.array.clone() - rhs.array.clone();
+        let array = array.into_shared();
+        let shape = lhs.shape.higher(&rhs.shape);
+
+        NdArrayTensor { array, shape }
+    }
+    fn sub_scalar<const D: usize>(
+        lhs: &<NdArrayBackend<E> as Backend>::TensorPrimitive<D>,
+        rhs: &E,
+    ) -> <NdArrayBackend<E> as Backend>::TensorPrimitive<D> {
+        let array = lhs.array.clone() - *rhs;
+        let shape = lhs.shape;
+
         NdArrayTensor { array, shape }
     }
 }
diff --git a/burn-tensor/src/tensor/backend/tch/ops/mod.rs b/burn-tensor/src/tensor/backend/tch/ops/mod.rs
@@ -16,5 +16,4 @@ mod neg;
 mod pow;
 mod precision;
 mod reshape;
-mod sub;
 mod transpose;