patch/0014-dts-rsqrt-nan.patch

From 9effaf54c572f7151396ff376168209c00b56b62 Mon Sep 17 00:00:00 2001
From: Tommy Chiang <ototot@google.com>
Date: Wed, 16 Oct 2024 22:37:09 -0700
Subject: [PATCH] rsqrt should report error when the input contains negative
 values

Due to the change of the behaviour, RsqrtNanInt{8,16} is now renamed to
RsqrtNegativeInt{8,16}

PiperOrigin-RevId: 686774011

PATCH_NAME=dts-rsqrt-nan
---
 tensorflow/lite/kernels/elementwise.cc      |  6 ++++++
 tensorflow/lite/kernels/elementwise_test.cc | 22 +++------------------
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index 5552f125..c71fe04d 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -384,6 +384,12 @@ TfLiteStatus RsqrtEvalQuantizedInt16(TfLiteContext* context, TfLiteNode* node,
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
   TfLiteTensor* output;
   TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  const int64_t num_elements = NumElements(input);
+  const int16_t* in_data = GetTensorData<int16_t>(input);
+  for (int64_t i = 0; i < num_elements; ++i) {
+    TF_LITE_ENSURE_MSG(context, in_data[i] >= op_data->input_offset,
+                       "Rsqrt is only defined for positive values");
+  }
   reference_integer_ops::LookupTable(
       GetTensorData<int16_t>(input),
       MatchingFlatSize(GetTensorShape(input), GetTensorShape(output)),
diff --git a/tensorflow/lite/kernels/elementwise_test.cc b/tensorflow/lite/kernels/elementwise_test.cc
index 97e42fb9..aa0965d2 100644
--- a/tensorflow/lite/kernels/elementwise_test.cc
+++ b/tensorflow/lite/kernels/elementwise_test.cc
@@ -417,13 +417,9 @@ TEST(ElementWise, RsqrtCloseTo0Int8) {
               ElementsAreArray(ArrayFloatNear(expected_output, kInputScale)));
 }
 
-TEST(ElementWise, RsqrtNanInt8) {
+TEST(ElementWise, RsqrtNegativeInt8) {
   const std::vector<float> input_data = {15., 46., 78.,  142.,
                                          1.,  17., -49., 113.};
-  std::vector<float> expected_output(input_data.size());
-  for (int i = 0; i < expected_output.size(); i++) {
-    expected_output[i] = 1.f / std::sqrt(input_data[i]);
-  }
   float kInputScale = 142.0 / 127.0;
   float kOutputScale = 1.0 / 255.0;
   int32_t input_zero_point = 0;
@@ -478,24 +474,12 @@ TEST(ElementWise, RsqrtInt16) {
       ElementsAreArray(ArrayFloatNear(expected_output, kQuantizedTolerance)));
 }
 
-TEST(ElementWise, RsqrtNanInt16) {
-  const float input_min = -0.8f;
-  const float input_max = 0.8f;
-
-  const float output_min = -2.4f;
-  const float output_max = 2.4f;
-
-  const float kQuantizedTolerance =
-      GetLUTTolerance<int16_t>(input_min, input_max, output_min, output_max);
-
+TEST(ElementWise, RsqrtNegativeInt16) {
   ElementWiseOpQuantizedModel m(BuiltinOperator_RSQRT,
                                 {TensorType_INT16, {1, 1, 4, 1}, -10, 10},
                                 {TensorType_INT16, {1, 1, 4, 1}, -10, 10});
   m.QuantizeAndPopulate<int16_t>(m.input(), {-1, 0, -4, -9});
-  ASSERT_EQ(m.Invoke(), kTfLiteOk);
-  EXPECT_THAT(m.ExtractDequantVector<int16_t>(m.output()),
-              ElementsAreArray(
-                  ArrayFloatNear({10, 9.82452, 10, 10}, kQuantizedTolerance)));
+  ASSERT_EQ(m.Invoke(), kTfLiteError);
 }
 
 TEST(ElementWise, Square) {