From d6442fd5490229565c806af3c37f0f8f48afb4f0 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Fri, 10 Jan 2025 08:46:46 -0800
Subject: [PATCH] Handle missing dtype cases in
 `xla::ifrt::DType::DebugString()`

PiperOrigin-RevId: 714066325
---
 xla/python/ifrt/dtype.cc | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/xla/python/ifrt/dtype.cc b/xla/python/ifrt/dtype.cc
index a79240f51a7e2..ed68a1d11403c 100644
--- a/xla/python/ifrt/dtype.cc
+++ b/xla/python/ifrt/dtype.cc
@@ -214,6 +214,10 @@ std::string DType::DebugString() const {
       return "INVALID";
     case kPred:
       return "PRED";
+    case kS2:
+      return "S2";
+    case kS4:
+      return "S4";
     case kS8:
       return "S8";
     case kS16:
@@ -222,6 +226,10 @@ std::string DType::DebugString() const {
       return "S32";
     case kS64:
       return "S64";
+    case kU2:
+      return "U2";
+    case kU4:
+      return "U4";
     case kU8:
       return "U8";
     case kU16:
@@ -246,6 +254,20 @@ std::string DType::DebugString() const {
       return "TOKEN";
     case kOpaque:
       return "OPAQUE";
+    case kF8E3M4:
+      return "F8E3M4";
+    case kF8E4M3:
+      return "F8E4M3";
+    case kF8E4M3FN:
+      return "F8E4M3FN";
+    case kF8E4M3B11FNUZ:
+      return "F8E4M3B11FNUZ";
+    case kF8E4M3FNUZ:
+      return "F8E4M3FNUZ";
+    case kF8E5M2:
+      return "F8E5M2";
+    case kF8E5M2FNUZ:
+      return "F8E5M2FNUZ";
     case kString:
       return "STRING";
     default: