[TFLite][Frontend] Support quantized SQUARE (#15914)

* add quantized square * add tests to quantized square
apache · Nov 9, 2023 · 99336c3 · 99336c3
1 parent 50597ce
commit 99336c3
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
@@ -1298,11 +1298,6 @@ def convert_square(self, op):
         assert len(output_tensors) == 1, "output tensors length should be 1"
         output_tensor = output_tensors[0]
 
-        if self.is_quantized(op):
-            raise tvm.error.OpNotImplemented(
-                "TFlite quantized SQUARE operator is not supported yet."
-            )
-
         exp_type = self.get_tensor_type_str(output_tensor.tensor.Type())
         out = _op.power(in_expr, relay.const(2, exp_type))
 

diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
@@ -2414,7 +2414,7 @@ def test_all_unary_elemwise():
     _test_forward_unary_elemwise(_test_floor)
     _test_forward_unary_elemwise(_test_exp)
     _test_forward_unary_elemwise(_test_log, negative=False)
-    _test_forward_unary_elemwise(_test_square)
+    _test_forward_unary_elemwise(_test_square, int_quant_dtype=tf.int8)
     _test_forward_unary_elemwise(_test_sin)
     _test_forward_unary_elemwise(_test_neg)
     _test_forward_unary_elemwise(_test_sqrt, negative=False)