Implement Element for bool

crates/burn-ndarray/src/element.rs

@@ -1,6 +1,5 @@
 use burn_tensor::Element;
 use ndarray::LinalgScalar;
-use num_traits::One;
 use num_traits::Signed;
 
 #[cfg(not(feature = "std"))]
@@ -20,7 +19,6 @@ where
 /// A general element for ndarray backend.
 pub trait NdArrayElement:
     Element
-    + One
     + ndarray::LinalgScalar
     + ndarray::ScalarOperand
     + ExpElement

crates/burn-ndarray/src/ops/base.rs

@@ -567,12 +567,12 @@ where
             tensor
                 .array
                 .mapv(|x| {
-                    if x > E::zero() {
-                        E::one()
-                    } else if x < E::zero() {
-                        -E::one()
+                    if x > <E as burn_tensor::identities::Zero>::zero() {
+                        <E as burn_tensor::identities::One>::one()
+                    } else if x < <E as burn_tensor::identities::Zero>::zero() {
+                        -<E as burn_tensor::identities::One>::one()
                     } else {
-                        E::zero()
+                        <E as burn_tensor::identities::Zero>::zero()
                     }
                 })
                 .into_shared(),

crates/burn-ndarray/src/ops/macros.rs

@@ -63,7 +63,11 @@ pub(crate) fn prod_dim<E: NdArrayElement, const D1: usize, const D2: usize>(
 ) -> NdArrayTensor<E, D2> {
     let array = tensor
         .array
-        .fold_axis(Axis(dim), E::one(), |acc, &x| acc.mul(x.elem()))
+        .fold_axis(
+            Axis(dim),
+            <E as burn_tensor::identities::One>::one(),
+            |acc, &x| acc.mul(x.elem()),
+        )
         .into_shared();
 
     NdArrayTensor { array }

crates/burn-tensor/src/tensor/api/numeric.rs

@@ -3,10 +3,9 @@ use alloc::vec::Vec;
 use crate::alloc::borrow::ToOwned;
 
 use crate::{
-    backend::Backend, check, check::TensorCheck, BasicOps, Bool, Distribution, Element,
-    ElementConversion, Float, Int, Shape, Tensor, TensorKind,
+    backend::Backend, check, check::TensorCheck, identities::Zero, BasicOps, Bool, Distribution,
+    Element, ElementConversion, Float, Int, Shape, Tensor, TensorKind,
 };
-use num_traits::Zero;
 
 impl<B, const D: usize, K> Tensor<B, D, K>
 where

crates/burn-tensor/src/tensor/data.rs

@@ -192,18 +192,6 @@ impl<E: Element> DataSerialize<E> {
     }
 }
 
-impl<const D: usize> Data<bool, D> {
-    /// Converts the data to a different element type.
-    pub fn convert<E: Element>(self) -> Data<E, D> {
-        let value: Vec<E> = self.value.into_iter().map(|a| (a as i64).elem()).collect();
-
-        Data {
-            value,
-            shape: self.shape,
-        }
-    }
-}
-
 impl<E: Element, const D: usize> Data<E, D> {
     /// Populates the data with random values.
     pub fn random<R: RngCore>(shape: Shape<D>, distribution: Distribution, rng: &mut R) -> Self {

crates/burn-tensor/src/tensor/element.rs → crates/burn-tensor/src/tensor/element/base.rs

@@ -1,8 +1,11 @@
 use core::cmp::Ordering;
 
-use crate::Distribution;
+use crate::{
+    cast::ToPrimitive,
+    identities::{One, Zero},
+    Distribution,
+};
 use half::{bf16, f16};
-use num_traits::{identities::Zero, One, ToPrimitive};
 use rand::RngCore;
 use serde::{Deserialize, Serialize};
 
@@ -223,6 +226,17 @@ make_element!(
     dtype DType::BF16
 );
 
+make_element!(
+    ty bool Precision::Other,
+    convert |elem: &dyn ToPrimitive| elem.to_u8().unwrap() != 0,
+    random |distribution: Distribution, rng: &mut R| {
+        let sample: u8 = distribution.sampler(rng).sample();
+        bool::from_elem(sample)
+    },
+    cmp |a: &bool, b: &bool| Ord::cmp(a, b),
+    dtype DType::Bool
+);
+
 #[allow(missing_docs)]
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub enum DType {