Make ElementComparison optional for dtypes

burn-book/src/advanced/backend-extension/README.md

@@ -34,7 +34,7 @@ pub trait Backend: burn::tensor::backend::Backend {
 You can then implement your new custom backend trait for any backend that you want to support:
 
 ```rust, ignore
-impl<E: TchElement> Backend for burn_tch::LibTorch<E> {
+impl<E: TchElement> Backend for burn_tch::LibTorch<E,F,I> {
    fn my_new_function(tensor: TchTensor<E, 2>) -> TchTensor<E, 2> {
       // My Tch implementation
    }
@@ -63,7 +63,7 @@ impl<B: Backend> Backend for burn_autodiff::Autodiff<B> {
    }
 }
 
-impl<E: TchElement> Backend for burn_autodiff::Autodiff<burn_tch::LibTorch<E>> {
+impl<E: TchElement> Backend for burn_autodiff::Autodiff<burn_tch::LibTorch<E,F,I>> {
    fn my_new_function(tensor: AutodiffTensor<E, 2>) -> AutodiffTensor<E, 2> {
       // My own backward implementation, generic over a backend implementation.
       //

crates/burn-backend/src/backend/base.rs

@@ -6,8 +6,8 @@ use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::element::Element;
-use crate::ops::*;
 use crate::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor};
+use crate::{ElementComparison, ops::*};
 use crate::{QTensorPrimitive, TensorData, TensorMetadata};
 
 use super::DeviceOps;
@@ -83,12 +83,12 @@ pub trait Backend:
     /// Tensor primitive to be used for all float operations.
     type FloatTensorPrimitive: TensorMetadata + 'static;
     /// Default float element type.
-    type FloatElem: Element;
+    type FloatElem: Element + ElementComparison;
 
     /// Tensor primitive to be used for all int operations.
     type IntTensorPrimitive: TensorMetadata + 'static;
     /// Int element type.
-    type IntElem: Element;
+    type IntElem: Element + ElementComparison;
 
     /// Tensor primitive to be used for all bool operations.
     type BoolTensorPrimitive: TensorMetadata + 'static;

crates/burn-backend/src/backend/ops/sort.rs

@@ -34,7 +34,7 @@ pub fn sort<B: Backend, K: TensorKind<B> + BasicOps<B>>(
     descending: bool,
 ) -> K::Primitive
 where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: ElementComparison,
 {
     let device = K::device(&tensor);
     let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation.";
@@ -51,7 +51,7 @@ pub fn sort_data<B: Backend, K: TensorKind<B> + BasicOps<B>>(
     descending: bool,
 ) -> K::Primitive
 where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: ElementComparison,
 {
     let dims = data.shape.clone();
     let data_slice = data.as_mut_slice().unwrap();
@@ -92,7 +92,7 @@ pub fn sort_with_indices<B: Backend, K: TensorKind<B> + BasicOps<B>>(
     descending: bool,
 ) -> (K::Primitive, IntTensor<B>)
 where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: ElementComparison,
 {
     let device = K::device(&tensor);
     let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation.";
@@ -109,7 +109,7 @@ fn sort_data_with_indices<B: Backend, K: TensorKind<B> + BasicOps<B>>(
     descending: bool,
 ) -> (K::Primitive, IntTensor<B>)
 where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: Element + ElementComparison,
 {
     let dims = data.shape.clone();
     let mut indices_data = dim_indices::<B>(&dims, dim);
@@ -191,7 +191,7 @@ pub fn argsort<B: Backend, K: TensorKind<B> + BasicOps<B>>(
     descending: bool,
 ) -> IntTensor<B>
 where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: ElementComparison,
 {
     let device = K::device(&tensor);
     let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation.";
@@ -209,7 +209,7 @@ fn argsort_data<B: Backend, K: TensorKind<B> + BasicOps<B>>(
     descending: bool,
 ) -> IntTensor<B>
 where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: ElementComparison,
 {
     let dims = data.shape.clone();
     let mut indices_data = dim_indices::<B>(&dims, dim);
@@ -252,7 +252,7 @@ fn sort_slice<B: Backend, K: BasicOps<B>>(
     permute_both: bool,
     descending: bool,
 ) where
-    <K as BasicOps<B>>::Elem: Element,
+    <K as BasicOps<B>>::Elem: ElementComparison,
 {
     let ndims = dims.len();
     let strides = compute_strides(dims);

crates/burn-backend/src/data/compare.rs

@@ -4,7 +4,7 @@ use burn_std::{DType, bf16, f16};
 use num_traits::{Float, ToPrimitive};
 
 use super::TensorData;
-use crate::element::Element;
+use crate::{ElementComparison, element::Element};
 
 /// The tolerance used to compare to floating point numbers.
 ///
@@ -269,7 +269,7 @@ impl TensorData {
         let mut num_diff = 0;
         let max_num_diff = 5;
         for (i, (a, b)) in self.iter::<E>().zip(other.iter::<E>()).enumerate() {
-            if a.cmp(&b).is_ne() {
+            if !a.eq(&b) {
                 // Only print the first 5 different values.
                 if num_diff < max_num_diff {
                     message += format!("\n  => Position {i}: {a} != {b}").as_str();
@@ -362,7 +362,7 @@ impl TensorData {
     ///
     /// If any value is not within the half-open range bounded inclusively below
     /// and exclusively above (`start..end`).
-    pub fn assert_within_range<E: Element>(&self, range: core::ops::Range<E>) {
+    pub fn assert_within_range<E: Element + ElementComparison>(&self, range: core::ops::Range<E>) {
         for elem in self.iter::<E>() {
             if elem.cmp(&range.start).is_lt() || elem.cmp(&range.end).is_ge() {
                 panic!("Element ({elem:?}) is not within range {range:?}");
@@ -379,7 +379,10 @@ impl TensorData {
     /// # Panics
     ///
     /// If any value is not within the half-open range bounded inclusively (`start..=end`).
-    pub fn assert_within_range_inclusive<E: Element>(&self, range: core::ops::RangeInclusive<E>) {
+    pub fn assert_within_range_inclusive<E: Element + ElementComparison>(
+        &self,
+        range: core::ops::RangeInclusive<E>,
+    ) {
         let start = range.start();
         let end = range.end();
 

crates/burn-backend/src/element/base.rs

@@ -14,7 +14,7 @@ pub trait Element:
     ToElement
     + ElementRandom
     + ElementConversion
-    + ElementComparison
+    + ElementEquality
     + ElementLimits
     + bytemuck::CheckedBitPattern
     + bytemuck::NoUninit
@@ -63,6 +63,12 @@ pub trait ElementRandom {
     fn random<R: RngCore>(distribution: Distribution, rng: &mut R) -> Self;
 }
 
+/// Element trait for equality of a tensor.
+pub trait ElementEquality {
+    /// Returns whether `self` and `other` are equal.
+    fn eq(&self, other: &Self) -> bool;
+}
+
 /// Element ordering trait.
 pub trait ElementComparison {
     /// Returns and [Ordering] between `self` and `other`.
@@ -104,6 +110,11 @@ macro_rules! make_element {
                 $dtype
             }
         }
+        impl ElementEquality for $type {
+            fn eq(&self, other: &Self) -> bool {
+                self == other
+            }
+        }
 
         impl ElementConversion for $type {
             #[inline(always)]