From f9a73552a6fbd47ebc30dc8f1919e1c25bf8a1a7 Mon Sep 17 00:00:00 2001
From: David Majnemer <majnemer@google.com>
Date: Mon, 5 Oct 2020 09:31:15 -0700
Subject: [PATCH] [XLA] Add NaN tests for RoundTripFpToString

Specifically, we want to check that NaN values with a sign bit set are appropriately converted.

PiperOrigin-RevId: 335433764
Change-Id: I2c6056645264054819b15e566e0489956147eb6e
---
 tensorflow/compiler/xla/BUILD        |  1 +
 tensorflow/compiler/xla/util_test.cc | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD
index cb202ce7600..c5233da6e9a 100644
--- a/tensorflow/compiler/xla/BUILD
+++ b/tensorflow/compiler/xla/BUILD
@@ -277,6 +277,7 @@ tf_cc_test(
         ":types",
         ":util",
         "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:bfloat16",
     ],
 )
 
diff --git a/tensorflow/compiler/xla/util_test.cc b/tensorflow/compiler/xla/util_test.cc
index 69acc59d8a2..5477dfba18d 100644
--- a/tensorflow/compiler/xla/util_test.cc
+++ b/tensorflow/compiler/xla/util_test.cc
@@ -15,10 +15,12 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/util.h"
 
+#include <limits>
 #include <list>
 
 #include "tensorflow/compiler/xla/test.h"
 #include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/core/platform/bfloat16.h"
 
 namespace xla {
 namespace {
@@ -103,5 +105,26 @@ TEST(UtilTest, SanitizeFileName) {
   EXPECT_EQ(SanitizeFileName("/A\\B[C]"), "_A_B_C_");
 }
 
+TEST(UtilTest, RoundTripFpToString) {
+  EXPECT_EQ(RoundTripFpToString(std::numeric_limits<Eigen::half>::quiet_NaN()),
+            "nan");
+  EXPECT_EQ(RoundTripFpToString(-std::numeric_limits<Eigen::half>::quiet_NaN()),
+            "-nan");
+  EXPECT_EQ(RoundTripFpToString(
+                std::numeric_limits<tensorflow::bfloat16>::quiet_NaN()),
+            "nan");
+  EXPECT_EQ(RoundTripFpToString(
+                -std::numeric_limits<tensorflow::bfloat16>::quiet_NaN()),
+            "-nan");
+  EXPECT_EQ(RoundTripFpToString(std::numeric_limits<float>::quiet_NaN()),
+            "nan");
+  EXPECT_EQ(RoundTripFpToString(-std::numeric_limits<float>::quiet_NaN()),
+            "-nan");
+  EXPECT_EQ(RoundTripFpToString(std::numeric_limits<double>::quiet_NaN()),
+            "nan");
+  EXPECT_EQ(RoundTripFpToString(-std::numeric_limits<double>::quiet_NaN()),
+            "-nan");
+}
+
 }  // namespace
 }  // namespace xla