Allow explicit casts to any type.

... as long as `UnderlyingTy` allows such a cast. XLA has some generic code that casts between types such as `std::complex` or `absl::int128` which are not marked as arithmetic types. PiperOrigin-RevId: 563481854
jax-ml · Sep 7, 2023 · 2ca30a2 · 2ca30a2
1 parent 780b6d0
commit 2ca30a2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/ml_dtypes/include/int4.h b/ml_dtypes/include/int4.h
@@ -48,7 +48,7 @@ struct i4 {
     return std::is_signed<UnderlyingTy>::value ? i4(7) : i4(15);
   }
 
-  template <typename T, typename = std::enable_if_t<std::is_arithmetic_v<T>>>
+  template <typename T>
   explicit constexpr operator T() const {
     return static_cast<T>(v);
   }

diff --git a/ml_dtypes/tests/int4_test.cc b/ml_dtypes/tests/int4_test.cc
@@ -287,12 +287,26 @@ TYPED_TEST(Int4Test, ToString) {
   }
 }
 
+struct CustomInt {
+  constexpr CustomInt() : x(0) {}
+  constexpr CustomInt(int x) : x(x) {}
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  constexpr operator int() const { return x; }
+  constexpr bool operator==(const CustomInt& other) const {
+    return x == other.x;
+  }
+
+ private:
+  int x;
+};
+
 #define GEN_DEST_TYPES(Type)                                               \
   std::pair<Type, bool>, std::pair<Type, uint4>, std::pair<Type, uint8_t>, \
       std::pair<Type, uint16_t>, std::pair<Type, uint32_t>,                \
       std::pair<Type, uint64_t>, std::pair<Type, int4>,                    \
       std::pair<Type, int8_t>, std::pair<Type, int16_t>,                   \
-      std::pair<Type, int32_t>, std::pair<Type, int64_t>
+      std::pair<Type, int32_t>, std::pair<Type, int64_t>,                  \
+      std::pair<Type, CustomInt>
 
 #define GEN_TYPE_PAIRS() GEN_DEST_TYPES(int4), GEN_DEST_TYPES(uint4)