Remove dtype rewrite (#2528)

* Remove dtype rewrite * Remove test * Fix tensor display for bool encoding abstraction --------- Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>
tracel-ai · Nov 25, 2024 · fe3e43a · fe3e43a
1 parent a35321c
commit fe3e43a
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 22 deletions.
diff --git a/crates/burn-candle/src/tensor.rs b/crates/burn-candle/src/tensor.rs
@@ -14,9 +14,7 @@ pub struct CandleTensor {
 impl TensorMetadata for CandleTensor {
     fn dtype(&self) -> DType {
         match self.tensor.dtype() {
-            // NOTE: bool tensors are stored as u32, we currently make this assumption
-            // since `TensorMetadata::dtype()` is used for display purposes only at this time.
-            candle_core::DType::U8 => DType::Bool,
+            candle_core::DType::U8 => DType::U8,
             candle_core::DType::U32 => DType::U32,
             candle_core::DType::I64 => DType::I64,
             candle_core::DType::BF16 => DType::BF16,

diff --git a/crates/burn-jit/src/tensor/base.rs b/crates/burn-jit/src/tensor/base.rs
@@ -65,12 +65,7 @@ where
 
 impl<R: JitRuntime> TensorMetadata for JitTensor<R> {
     fn dtype(&self) -> DType {
-        match self.dtype {
-            // NOTE: bool tensors are stored as u32, we currently make this assumption
-            // since `TensorMetadata::dtype()` is used for display purposes only at this time.
-            DType::U32 => DType::Bool,
-            _ => self.dtype,
-        }
+        self.dtype
     }
 
     fn shape(&self) -> Shape {

diff --git a/crates/burn-tensor/src/tensor/api/base.rs b/crates/burn-tensor/src/tensor/api/base.rs
@@ -7,6 +7,7 @@ use alloc::string::String;
 use alloc::vec;
 
 use burn_common::stub::RwLock;
+use core::any::TypeId;
 use core::future::Future;
 use core::iter::repeat;
 use core::{fmt::Debug, ops::Range};
@@ -1873,7 +1874,15 @@ where
         writeln!(f, "  device:  {:?},", self.device())?;
         writeln!(f, "  backend:  {:?},", B::name())?;
         writeln!(f, "  kind:  {:?},", K::name())?;
-        writeln!(f, "  dtype:  {:?},", self.primitive.dtype().name())?;
+
+        // Bool tensors might be encoded in a different type, which we abstract for the display
+        let dtype = if TypeId::of::<K::Elem>() == TypeId::of::<bool>() {
+            DType::Bool
+        } else {
+            self.primitive.dtype()
+        };
+
+        writeln!(f, "  dtype:  {:?},", dtype.name())?;
         write!(f, "}}")
     }
 }

diff --git a/crates/burn-tensor/src/tests/primitive.rs b/crates/burn-tensor/src/tests/primitive.rs
@@ -30,16 +30,4 @@ mod tests {
             <TestBackend as Backend>::IntElem::dtype() // default int elem type
         );
     }
-
-    #[test]
-    fn should_support_bool_dtype() {
-        let tensor =
-            TestTensorBool::<2>::from([[false, true, true], [false, false, true]]).into_primitive();
-
-        assert_eq!(
-            burn_tensor::TensorMetadata::shape(&tensor),
-            Shape::new([2, 3])
-        );
-        assert_eq!(burn_tensor::TensorMetadata::dtype(&tensor), DType::Bool);
-    }
 }