Merge pull request #395 from vaijira/add_neg

Add Neg trait to Float trait.
tracel-ai · Jan 6, 2025 · 4d6f50f · 4d6f50f
2 parents 63b27a3 + 05aebdb
commit 4d6f50f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/crates/cubecl-core/src/frontend/element/float.rs b/crates/cubecl-core/src/frontend/element/float.rs
@@ -35,6 +35,7 @@ pub trait Float:
     + Normalize
     + Dot
     + Into<Self::ExpandType>
+    + core::ops::Neg<Output = Self>
     + core::ops::Add<Output = Self>
     + core::ops::Sub<Output = Self>
     + core::ops::Mul<Output = Self>

diff --git a/crates/cubecl-core/src/frontend/element/float/tensor_float.rs b/crates/cubecl-core/src/frontend/element/float/tensor_float.rs
@@ -2,7 +2,7 @@
 #![allow(clippy::transmute_float_to_int)] // prev=1.83.
 
 use bytemuck::{Pod, Zeroable};
-use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};
+use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
 use half::f16;
 use num_traits::{NumCast, ToPrimitive};
 use serde::Serialize;
@@ -88,6 +88,14 @@ impl tf32 {
     }
 }
 
+impl Neg for tf32 {
+    type Output = Self;
+
+    fn neg(self) -> Self::Output {
+        Self::from_f32(self.to_f32().neg())
+    }
+}
+
 impl Mul for tf32 {
     type Output = Self;