Merge pull request #202 from sergey-kozub:signbit

PiperOrigin-RevId: 675612039
jax-ml · Sep 17, 2024 · 620f804 · 620f804
2 parents b65a1f6 + 0246565
commit 620f804
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/ml_dtypes/include/mxfloat.h b/ml_dtypes/include/mxfloat.h
@@ -338,7 +338,7 @@ struct Traits<float4_e2m1fn>
 namespace Eigen {
 namespace numext {
 
-#define MXFLOAT_EIGEN_BITCAST_IMPL(Type)                                 \
+#define MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(Type)                     \
   template <>                                                            \
   EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint8_t bit_cast<uint8_t, Type>( \
       const Type& x) {                                                   \
@@ -348,11 +348,16 @@ namespace numext {
   EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Type bit_cast<Type, uint8_t>(    \
       const uint8_t& x) {                                                \
     return Type::FromRep(x);                                             \
+  }                                                                      \
+  template <>                                                            \
+  EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Type signbit(const Type& x) {    \
+    int8_t t = bit_cast<int8_t, Type>(x) << (8 - Type::kBits);           \
+    return bit_cast<Type, int8_t>(t >> 7);                               \
   }
 
-MXFLOAT_EIGEN_BITCAST_IMPL(ml_dtypes::float6_e2m3fn)
-MXFLOAT_EIGEN_BITCAST_IMPL(ml_dtypes::float6_e3m2fn)
-MXFLOAT_EIGEN_BITCAST_IMPL(ml_dtypes::float4_e2m1fn)
+MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(ml_dtypes::float6_e2m3fn)
+MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(ml_dtypes::float6_e3m2fn)
+MXFLOAT_EIGEN_BITCAST_AND_SIGNBIT_IMPL(ml_dtypes::float4_e2m1fn)
 
 #undef MXFLOAT_EIGEN_BITCAST_IMPL
 

diff --git a/ml_dtypes/tests/mxfloat_test.cc b/ml_dtypes/tests/mxfloat_test.cc
@@ -125,6 +125,14 @@ TYPED_TEST(FloatMXTest, Negate) {
   }
 }
 
+TYPED_TEST(FloatMXTest, Signbit) {
+  using FloatMX = TypeParam;
+
+  FloatMX one(1.0);
+  EXPECT_EQ(Eigen::numext::signbit(one).rep(), 0x00);
+  EXPECT_EQ(Eigen::numext::signbit(-one).rep(), 0xff);
+}
+
 TYPED_TEST(FloatMXTest, BitCasts) {
   using FloatMX = TypeParam;