[XLA] Add S1 and U1 as data types

PiperOrigin-RevId: 706011071
jax-ml · Dec 14, 2024 · 86b812e · 86b812e
1 parent 401ed6a
commit 86b812e
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/ml_dtypes/include/intn.h b/ml_dtypes/include/intn.h
@@ -241,7 +241,9 @@ struct intN {
   }
 };
 
+using int1 = intN<1, int8_t>;
 using int2 = intN<2, int8_t>;
+using uint1 = intN<1, uint8_t>;
 using uint2 = intN<2, uint8_t>;
 using int4 = intN<4, int8_t>;
 using uint4 = intN<4, uint8_t>;
@@ -294,7 +296,12 @@ struct intN_numeric_limits_base {
 }  // namespace ml_dtypes
 
 namespace std {
-
+template <>
+struct numeric_limits<ml_dtypes::int1>
+    : public ml_dtypes::internal::intN_numeric_limits_base<ml_dtypes::int1> {};
+template <>
+struct numeric_limits<ml_dtypes::uint1>
+    : public ml_dtypes::internal::intN_numeric_limits_base<ml_dtypes::uint1> {};
 template <>
 struct numeric_limits<ml_dtypes::int2>
     : public ml_dtypes::internal::intN_numeric_limits_base<ml_dtypes::int2> {};