feat: add comparison (tracel-ai#19)

djdisodo · Aug 22, 2022 · f8ab29b · f8ab29b
1 parent 5560bae
commit f8ab29b
Show file tree

Hide file tree

Showing 8 changed files with 155 additions and 3 deletions.
diff --git a/.github/workflows/test-burn-dataset.yml b/.github/workflows/test-burn-dataset.yml
@@ -0,0 +1,29 @@
+name: test
+
+on: [push]
+
+jobs:
+  publish:
+    name: test burn dataset
+    runs-on: ubuntu-latest
+    steps:
+      - name: checkout
+        uses: actions/checkout@v2
+
+      - name: install rust nightly
+        uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: nightly
+          components: rustfmt
+          override: true
+
+      - name: check format
+        run: |
+          cd burn-dataset
+          cargo fmt --check --all
+
+      - name: check tests
+        run: |
+          cd burn-dataset
+          cargo test
diff --git a/.github/workflows/test-burn.yml b/.github/workflows/test-burn.yml
@@ -0,0 +1,27 @@
+name: test
+
+on: [push]
+
+jobs:
+  publish:
+    name: test burn
+    runs-on: ubuntu-latest
+    steps:
+      - name: checkout
+        uses: actions/checkout@v2
+
+      - name: install rust nightly
+        uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: nightly
+          components: rustfmt
+          override: true
+
+      - name: check format
+        run: |
+          cargo fmt --check --all
+
+      - name: check tests
+        run: |
+          cargo test
diff --git a/burn-tensor/src/tensor/api/tensor.rs b/burn-tensor/src/tensor/api/tensor.rs
@@ -159,6 +159,14 @@ where
         Self::new(self.value.sum_dim(dim))
     }
 
+    pub fn equal(&self, other: &Self) -> BoolTensor<B, D> {
+        BoolTensor::new(self.value.equal(&other.value))
+    }
+
+    pub fn equal_scalar(&self, other: &B::Elem) -> BoolTensor<B, D> {
+        BoolTensor::new(self.value.equal_scalar(other))
+    }
+
     pub fn greater(&self, other: &Self) -> BoolTensor<B, D> {
         BoolTensor::new(self.value.greater(&other.value))
     }

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/map_comparison.rs b/burn-tensor/src/tensor/backend/autodiff/ops/map_comparison.rs
@@ -10,6 +10,17 @@ macro_rules! define_impl {
         impl<E: $element, const D: usize> TensorOpsMapComparison<$backend, D>
             for <$backend as Backend>::TensorPrimitive<D>
         {
+            fn equal(&self, other: &Self) -> <$backend as Backend>::BoolTensorPrimitive<D> {
+                TensorOpsMapComparison::equal(&self.tensor(), &other.tensor())
+            }
+
+            fn equal_scalar(
+                &self,
+                other: &<$backend as Backend>::Elem,
+            ) -> <$backend as Backend>::BoolTensorPrimitive<D> {
+                TensorOpsMapComparison::equal_scalar(&self.tensor(), other)
+            }
+
             fn greater(&self, other: &Self) -> <$backend as Backend>::BoolTensorPrimitive<D> {
                 TensorOpsMapComparison::greater(&self.tensor(), &other.tensor())
             }

diff --git a/burn-tensor/src/tensor/backend/ndarray/ops/map_comparison.rs b/burn-tensor/src/tensor/backend/ndarray/ops/map_comparison.rs
@@ -6,6 +6,27 @@ impl<E, const D: usize> TensorOpsMapComparison<NdArrayBackend<E>, D> for NdArray
 where
     E: NdArrayElement,
 {
+    fn equal(
+        &self,
+        other: &Self,
+    ) -> <NdArrayBackend<E> as crate::back::Backend>::BoolTensorPrimitive<D> {
+        let tensor = self.sub(other);
+        let zero = E::zeros(&E::default());
+        tensor.equal_scalar(&zero)
+    }
+
+    fn equal_scalar(
+        &self,
+        other: &<NdArrayBackend<E> as crate::back::Backend>::Elem,
+    ) -> <NdArrayBackend<E> as crate::back::Backend>::BoolTensorPrimitive<D> {
+        let array = self.array.mapv(|a| a == *other).into_shared();
+
+        NdArrayTensor {
+            shape: self.shape,
+            array,
+        }
+    }
+
     fn greater(
         &self,
         other: &Self,

diff --git a/burn-tensor/src/tensor/backend/tch/ops/map_comparison.rs b/burn-tensor/src/tensor/backend/tch/ops/map_comparison.rs
@@ -9,6 +9,33 @@ impl<E, const D: usize> TensorOpsMapComparison<TchBackend<E>, D> for TchTensor<E
 where
     E: TchElement,
 {
+    fn equal(
+        &self,
+        other: &Self,
+    ) -> <TchBackend<E> as crate::back::Backend>::BoolTensorPrimitive<D> {
+        let tensor = self.tensor.eq_tensor(&other.tensor);
+
+        TchTensor {
+            shape: self.shape,
+            tensor,
+            kind: TchKind::<bool>::new(),
+        }
+    }
+
+    fn equal_scalar(
+        &self,
+        other: &<TchBackend<E> as crate::back::Backend>::Elem,
+    ) -> <TchBackend<E> as crate::back::Backend>::BoolTensorPrimitive<D> {
+        let other: f64 = (*other).into();
+        let tensor = self.tensor.eq(other);
+
+        TchTensor {
+            shape: self.shape,
+            tensor,
+            kind: TchKind::<bool>::new(),
+        }
+    }
+
     fn greater(
         &self,
         other: &Self,

diff --git a/burn-tensor/src/tensor/ops.rs b/burn-tensor/src/tensor/ops.rs
@@ -58,6 +58,8 @@ pub trait TensorOpsIndex<E, const D1: usize> {
 }
 
 pub trait TensorOpsMapComparison<B: Backend, const D: usize> {
+    fn equal(&self, other: &Self) -> B::BoolTensorPrimitive<D>;
+    fn equal_scalar(&self, other: &B::Elem) -> B::BoolTensorPrimitive<D>;
     fn greater(&self, other: &Self) -> B::BoolTensorPrimitive<D>;
     fn greater_scalar(&self, other: &B::Elem) -> B::BoolTensorPrimitive<D>;
     fn greater_equal(&self, other: &Self) -> B::BoolTensorPrimitive<D>;

diff --git a/burn-tensor/src/tensor/tensor_trait.rs b/burn-tensor/src/tensor/tensor_trait.rs
@@ -1,13 +1,14 @@
 use crate::{tensor::ops::*, Distribution};
 use half::f16;
-use num_traits::{One, ToPrimitive, Zero};
+use num_traits::ToPrimitive;
 use rand::prelude::StdRng;
 
 pub trait Element:
     Zeros<Self>
     + ToPrimitive
     + ElementRandom<Self>
     + ElementConversion
+    + ElementValue
     + Ones<Self>
     + std::ops::Mul<Self, Output = Self>
     + std::fmt::Debug
@@ -37,6 +38,14 @@ pub trait ElementRandom<T> {
     fn random(distribution: Distribution<T>, rng: &mut StdRng) -> T;
 }
 
+pub trait ElementValue {
+    fn inf() -> Self;
+    fn inf_neg() -> Self;
+    fn nan() -> Self;
+    fn zero() -> Self;
+    fn one() -> Self;
+}
+
 #[cfg(feature = "ndarray")]
 pub trait NdArrayElement:
     Element + ndarray::LinalgScalar + ndarray::ScalarOperand + ExpElement + num_traits::FromPrimitive
@@ -87,6 +96,24 @@ macro_rules! ad_items {
             }
         }
 
+        impl ElementValue for $float {
+            fn inf() -> Self {
+                Self::from_elem(f64::INFINITY)
+            }
+            fn inf_neg() -> Self {
+                Self::from_elem(std::ops::Neg::neg(f64::INFINITY))
+            }
+            fn nan() -> Self {
+                Self::from_elem(f64::NAN)
+            }
+            fn zero() -> Self {
+                $zero
+            }
+            fn one() -> Self {
+                $one
+            }
+        }
+
         impl ElementRandom<$float> for $float {
             fn random(distribution: Distribution<$float>, rng: &mut StdRng) -> $float {
                 $random(distribution, rng)
@@ -135,8 +162,8 @@ ad_items!(
 
 ad_items!(
     ty f16,
-    zero f16::zero(),
-    one f16::one(),
+    zero <f16 as num_traits::Zero>::zero(),
+    one <f16 as num_traits::One>::one(),
     convert |elem: &dyn ToPrimitive| f16::from_f32(elem.to_f32().unwrap()),
     random |distribution: Distribution<f16>, rng: &mut StdRng| {
         let distribution: Distribution<f32> = distribution.convert();