update tensorflow/python/

ROCm · Oct 1, 2024 · 069a82c · 069a82c
1 parent 1086598
commit 069a82c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/tensorflow/python/lib/core/ndarray_tensor.cc b/tensorflow/python/lib/core/ndarray_tensor.cc
@@ -216,7 +216,13 @@ Status PyArray_TYPE_to_TF_DataType(PyArrayObject* array,
       } else if (pyarray_type == custom_dtypes.float8_e4m3fn) {
         *out_tf_datatype = TF_FLOAT8_E4M3FN;
         break;
-      } else if (pyarray_type == custom_dtypes.int4) {
+      } else if (pyarray_type == custom_dtypes.float8_e5m2fnuz) {
+        *out_tf_datatype = TF_FLOAT8_E5M2FNUZ;
+        break;
+      } else if (pyarray_type == custom_dtypes.float8_e4m3fnuz) {
+        *out_tf_datatype = TF_FLOAT8_E4M3FNUZ;
+        break;
+      }else if (pyarray_type == custom_dtypes.int4) {
         *out_tf_datatype = TF_INT4;
         break;
       } else if (pyarray_type == custom_dtypes.uint4) {

diff --git a/tensorflow/python/lib/core/ndarray_tensor_bridge.cc b/tensorflow/python/lib/core/ndarray_tensor_bridge.cc
@@ -198,6 +198,12 @@ Status TF_DataType_to_PyArray_TYPE(TF_DataType tf_datatype,
     case TF_FLOAT8_E4M3FN:
       *out_pyarray_type = custom_dtypes.float8_e4m3fn;
       break;
+    case TF_FLOAT8_E5M2FNUZ:
+      *out_pyarray_type = custom_dtypes.float8_e5m2fnuz;
+      break;
+    case TF_FLOAT8_E4M3FNUZ:
+      *out_pyarray_type = custom_dtypes.float8_e4m3fnuz;
+      break;
     case TF_INT4:
       *out_pyarray_type = custom_dtypes.int4;
       break;