From f6720ec3d3daf06764aef39a4124241445bf7ec2 Mon Sep 17 00:00:00 2001
From: Ophir Frish <ophir.frish@arm.com>
Date: Wed, 27 Oct 2021 22:26:45 +0300
Subject: [PATCH] Support quantized NEG operator in TFLite frontend

---
 python/tvm/relay/frontend/tflite.py          |  2 -
 tests/python/frontend/tflite/test_forward.py | 51 +++++++++++++++-----
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
index 3688ff5ff4e5..d6148d31e148 100644
--- a/python/tvm/relay/frontend/tflite.py
+++ b/python/tvm/relay/frontend/tflite.py
@@ -1203,8 +1203,6 @@ def convert_rsqrt(self, op):
 
     def convert_neg(self, op):
         """Convert TFLite NEG"""
-        if self.is_quantized(op):
-            raise tvm.error.OpNotImplemented("TFlite quantized NEG operator is not supported yet.")
         return self._convert_unary_elemwise(_op.negative, op)
 
     def convert_elu(self, op):
diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
index 754976ca8c13..4ec4603a5606 100644
--- a/tests/python/frontend/tflite/test_forward.py
+++ b/tests/python/frontend/tflite/test_forward.py
@@ -1869,16 +1869,6 @@ def _test_sqrt(data):
     return _test_unary_elemwise(math_ops.sqrt, data)
 
 
-#######################################################################
-# Neg
-# ---
-
-
-def _test_neg(data):
-    """One iteration of neg"""
-    return _test_unary_elemwise(math_ops.neg, data)
-
-
 #######################################################################
 # Square
 # ------
@@ -1914,7 +1904,6 @@ def test_all_unary_elemwise():
     _test_forward_unary_elemwise(_test_log)
     _test_forward_unary_elemwise(_test_sin)
     _test_forward_unary_elemwise(_test_sqrt)
-    _test_forward_unary_elemwise(_test_neg)
     _test_forward_unary_elemwise(_test_square)
     # ceil and cos come with TFLite 1.14.0.post1 fbs schema
     if package_version.parse(tf.VERSION) >= package_version.parse("1.14.0"):
@@ -3381,6 +3370,45 @@ def test_forward_rsqrt():
     _test_rsqrt(np.arange(1, 240, 40, dtype=np.uint8).reshape((2, 1, 3)), quantized=True)
 
 
+#######################################################################
+# NEG
+# ----
+
+
+def _test_neg(data, quantized=False):
+    """One iteration of NEG"""
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=data.shape, dtype="float32", name="in_0")
+
+        if quantized:
+            inq_data = tf.quantization.fake_quant_with_min_max_args(
+                in_data, min=1, max=6, name="inq_0"
+            )
+            input_range = {"inq_0": (1, 6)}
+            out = math_ops.neg(inq_data)
+            out = tf.quantization.fake_quant_with_min_max_args(out, min=1, max=6, name="out")
+            compare_tflite_with_tvm(
+                data,
+                "inq_0:0",
+                [inq_data],
+                [out],
+                quantized=True,
+                input_range=input_range,
+                experimental_new_converter=True,
+            )
+        else:
+            out = math_ops.neg(in_data)
+            compare_tflite_with_tvm(data, "in_0:0", [in_data], [out])
+
+
+def test_forward_neg():
+    """NEG"""
+    _test_neg(np.arange(-2.0, 4.0, dtype=np.float32), quantized=False)
+    _test_neg(np.arange(-2.0, 4.0, dtype=np.float32).reshape((2, 1, 3)), quantized=False)
+    _test_neg(np.arange(1, 240, 40, dtype=np.uint8), quantized=True)
+    _test_neg(np.arange(1, 240, 40, dtype=np.uint8).reshape((2, 1, 3)), quantized=True)
+
+
 #######################################################################
 # ReLu
 # ----
@@ -4657,6 +4685,7 @@ def test_prevent_tensorflow_dynamic_range():
     test_forward_softmax()
     test_forward_tanh()
     test_forward_rsqrt()
+    test_forward_neg()
     test_forward_relu()
     test_forward_relu6()
     test_forward_leaky_relu()