Add Float4e2m1fn support in tensorstore.

PiperOrigin-RevId: 836265256
Change-Id: Id4c0848b7207420cd8fbe2f6856047cfc256b5d3

Author: Tensorstore Team

docs/tensorstore_schema.yml

@@ -115,6 +115,10 @@ definitions:
         title: |
           `8-bit floating point - Exponent: 5, Mantissa: 2, bias: 16, with NaN, without infinities.
           <https://github.com/jax-ml/ml_dtypes#float8_e5m2fnuz>`__ .
+      - const: "float4_e2m1fn"
+        title: |
+          `4-bit floating point - Exponent: 2, Mantissa: 1, bias: 1, without NaN and infinities.
+          <https://github.com/jax-ml/ml_dtypes#float4_e2m1fn>`__ .
       - const: "float16"
         title: |
           `IEEE 754 binary16

python/tensorstore/__init__.py

@@ -257,6 +257,15 @@ class Indexable(metaclass=abc.ABCMeta):
   Data types
 """
 
+float4_e2m1fn: dtype
+"""8-bit floating-point data type.
+
+Details in https://github.com/jax-ml/ml_dtypes#float4_e2m1fn
+
+Group:
+  Data types
+"""
+
 float16: dtype
 """:wikipedia:`IEEE 754 binary16 <Half-precision_floating-point_format>` half-precision floating-point data type.  Correspond to ``numpy.float16``.
 

python/tensorstore/__init__.pyi

@@ -76,6 +76,8 @@ float8_e5m2: dtype
 "8-bit floating-point data type.\n\nDetails in https://github.com/jax-ml/ml_dtypes#float8_e5m2\n\nGroup:\n  Data types\n"
 float8_e5m2fnuz: dtype
 "8-bit floating-point data type.\n\nDetails in https://github.com/jax-ml/ml_dtypes#float8_e5m2fnuz\n\nGroup:\n  Data types\n"
+float4_e2m1fn: dtype
+"8-bit floating-point data type.\n\nDetails in https://github.com/jax-ml/ml_dtypes#float4_e2m1fn\n\nGroup:\n  Data types\n"
 float16: dtype
 ":wikipedia:`IEEE 754 binary16 <Half-precision_floating-point_format>` half-precision floating-point data type.  Correspond to ``numpy.float16``.\n\nGroup:\n  Data types\n"
 bfloat16: dtype
@@ -169,6 +171,7 @@ __all__ = [
     "experimental_update_verbose_logging",
     "float16",
     "float32",
+    "float4_e2m1fn",
     "float64",
     "float8_e3m4",
     "float8_e4m3b11fnuz",

python/tensorstore/data_type.cc

@@ -85,6 +85,7 @@ class CustomDTypes {
   }                                                                     \
   /**/
     TENSORSTORE_FOR_EACH_FLOAT8_DATA_TYPE(TENSORSTORE_INTERNAL_DO_ADD_TO_MAP)
+    TENSORSTORE_FOR_EACH_MXFLOAT_DATA_TYPE(TENSORSTORE_INTERNAL_DO_ADD_TO_MAP)
     TENSORSTORE_FOR_EACH_LOW_PRECISION_INT_DATA_TYPE(
         TENSORSTORE_INTERNAL_DO_ADD_TO_MAP)
 #undef TENSORSTORE_INTERNAL_DO_ADD_TO_MAP
@@ -151,6 +152,8 @@ int GetNumpyTypeNum(DataType dtype) {
       TENSORSTORE_INTERNAL_DO_GET_NPY_TYPE_NUM_CASE(bfloat16_t)
       TENSORSTORE_FOR_EACH_FLOAT8_DATA_TYPE(
           TENSORSTORE_INTERNAL_DO_GET_NPY_TYPE_NUM_CASE)
+      TENSORSTORE_FOR_EACH_MXFLOAT_DATA_TYPE(
+          TENSORSTORE_INTERNAL_DO_GET_NPY_TYPE_NUM_CASE)
       TENSORSTORE_FOR_EACH_LOW_PRECISION_INT_DATA_TYPE(
           TENSORSTORE_INTERNAL_DO_GET_NPY_TYPE_NUM_CASE)
 #undef TENSORSTORE_INTERNAL_DO_GET_NPY_TYPE_NUM_CASE

python/tensorstore/tests/custom_dtypes_test.py

@@ -28,6 +28,7 @@
     ml_dtypes.float8_e4m3b11fnuz.dtype,
     ml_dtypes.float8_e5m2.dtype,
     ml_dtypes.float8_e5m2fnuz.dtype,
+    ml_dtypes.float4_e2m1fn.dtype,
     ml_dtypes.bfloat16.dtype,
     ts.int4.numpy_dtype,
     ts.float8_e3m4.numpy_dtype,
@@ -36,6 +37,7 @@
     ts.float8_e4m3b11fnuz.numpy_dtype,
     ts.float8_e5m2.numpy_dtype,
     ts.float8_e5m2fnuz.numpy_dtype,
+    ts.float4_e2m1fn.numpy_dtype,
     ts.bfloat16.numpy_dtype,
     np.dtype("int4"),
     np.dtype("float8_e3m4"),
@@ -44,6 +46,7 @@
     np.dtype("float8_e4m3b11fnuz"),
     np.dtype("float8_e5m2"),
     np.dtype("float8_e5m2fnuz"),
+    np.dtype("float4_e2m1fn"),
     np.dtype("bfloat16"),
 ]
 

tensorstore/BUILD

@@ -7,6 +7,8 @@ licenses(["notice"])
 
 exports_files(["LICENSE"])
 
+exports_files(["generate_data_type.py"])
+
 tensorstore_cc_library(
     name = "array",
     srcs = [
@@ -443,6 +445,7 @@ tensorstore_cc_library(
         "//tensorstore/util:float8",
         "//tensorstore/util:int2",
         "//tensorstore/util:int4",
+        "//tensorstore/util:mxfloat",
         "//tensorstore/util:result",
         "//tensorstore/util:str_cat",
         "//tensorstore/util:utf8_string",

tensorstore/data_type.cc

@@ -418,6 +418,9 @@ TENSORSTORE_INTERNAL_INHERITED_CONVERT(  //
     ::tensorstore::dtypes::complex64_t,
     ::tensorstore::dtypes::float8_e5m2fnuz_t,
     internal_data_type::ComplexNumericConvertDataType)
+TENSORSTORE_INTERNAL_INHERITED_CONVERT(  //
+    ::tensorstore::dtypes::complex64_t, ::tensorstore::dtypes::float4_e2m1fn_t,
+    internal_data_type::ComplexNumericConvertDataType)
 TENSORSTORE_INTERNAL_INHERITED_CONVERT(  //
     ::tensorstore::dtypes::complex64_t, ::tensorstore::dtypes::float16_t,
     internal_data_type::ComplexNumericConvertDataType)
@@ -451,6 +454,9 @@ TENSORSTORE_INTERNAL_INHERITED_CONVERT(  //
     ::tensorstore::dtypes::complex128_t,
     ::tensorstore::dtypes::float8_e5m2fnuz_t,
     internal_data_type::ComplexNumericConvertDataType)
+TENSORSTORE_INTERNAL_INHERITED_CONVERT(  //
+    ::tensorstore::dtypes::complex128_t, ::tensorstore::dtypes::float4_e2m1fn_t,
+    internal_data_type::ComplexNumericConvertDataType)
 TENSORSTORE_INTERNAL_INHERITED_CONVERT(  //
     ::tensorstore::dtypes::complex128_t, ::tensorstore::dtypes::float16_t,
     internal_data_type::ComplexNumericConvertDataType)

tensorstore/data_type.h

@@ -83,6 +83,7 @@
 #include "tensorstore/util/float8.h"
 #include "tensorstore/util/int2.h"
 #include "tensorstore/util/int4.h"
+#include "tensorstore/util/mxfloat.h"
 #include "tensorstore/util/str_cat.h"
 #include "tensorstore/util/utf8_string.h"
 
@@ -171,6 +172,9 @@ using float8_e5m2_t = ::tensorstore::Float8e5m2;
 using float8_e5m2fnuz_t = ::tensorstore::Float8e5m2fnuz;
 ///
 /// \ingroup data types
+using float4_e2m1fn_t = ::tensorstore::Float4e2m1fn;
+///
+/// \ingroup data types
 using bfloat16_t = ::tensorstore::BFloat16;
 
 /// :wikipedia:`IEEE 754 binary16<Half-precision_floating-point_format>`
@@ -243,6 +247,7 @@ enum class DataTypeId {
   float8_e4m3b11fnuz_t,
   float8_e5m2_t,
   float8_e5m2fnuz_t,
+  float4_e2m1fn_t,
   float16_t,
   bfloat16_t,
   float32_t,
@@ -295,8 +300,13 @@ inline constexpr size_t kNumDataTypeIds =
   X(float8_e5m2fnuz_t, ##__VA_ARGS__)                 \
   /**/
 
+#define TENSORSTORE_FOR_EACH_MXFLOAT_DATA_TYPE(X, ...) \
+  X(float4_e2m1fn_t, ##__VA_ARGS__)                    \
+  /**/
+
 #define TENSORSTORE_FOR_EACH_LOW_PRECISION_FLOAT_DATA_TYPE(X, ...) \
   TENSORSTORE_FOR_EACH_FLOAT8_DATA_TYPE(X, ##__VA_ARGS__)          \
+  TENSORSTORE_FOR_EACH_MXFLOAT_DATA_TYPE(X, ##__VA_ARGS__)         \
   X(float16_t, ##__VA_ARGS__)                                      \
   X(bfloat16_t, ##__VA_ARGS__)                                     \
   /**/