From 9a263a8c9b2f2dd6124cdeb26f28b5285b84d3d6 Mon Sep 17 00:00:00 2001 From: Dmitriy Smirnov Date: Tue, 24 Nov 2020 22:40:48 +0000 Subject: [PATCH] [TFLite] added scalar axis value handling in reduce Axis value in reduce can now be specified as scalar --- python/tvm/relay/frontend/tflite.py | 3 ++- tests/python/frontend/tflite/test_forward.py | 18 +++++++++++------- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index 623aeee358a6..6509c41b0821 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -1638,7 +1638,8 @@ def _convert_reduce(self, relay_op, op): in_expr = self.get_expr(input_tensor.tensor_idx) # axis - axis = tuple(self.get_tensor_value(input_tensors[1])) + axis_value = self.get_tensor_value(input_tensors[1]) + axis = tuple(axis_value) if len(axis_value.shape) > 0 else tuple((axis_value.item(),)) # Options - keep_dims (bool) assert op.BuiltinOptionsType() == BuiltinOptions.ReducerOptions diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py index b7f3b91f4243..d8fcf7081b3c 100644 --- a/tests/python/frontend/tflite/test_forward.py +++ b/tests/python/frontend/tflite/test_forward.py @@ -2143,18 +2143,22 @@ def _test_forward_reduce(testop, dtype="float32"): if dtype == "bool": data0 = [np.random.choice(a=[False, True], size=(16, 16, 16, 16)).astype(dtype), None] data1 = [ + np.random.choice(a=[False, True], size=(16, 16, 16, 16)).astype(dtype), + np.array(1, dtype=np.int32), + ] + data2 = [ np.random.choice(a=[False, True], size=(16, 16, 16, 16)).astype(dtype), np.array([1, 2], dtype=np.int32), ] else: data0 = [np.random.rand(16, 16, 16, 16).astype(dtype), None] - data1 = [np.random.rand(16, 16, 16, 16).astype(dtype), np.array([1, 2], dtype=np.int32)] - testop(data0) - testop(data0, keep_dims=False) - testop(data0, keep_dims=True) - testop(data1) - testop(data1, keep_dims=False) - testop(data1, keep_dims=True) + data1 = [np.random.rand(16, 16, 16, 16).astype(dtype), np.array(1, dtype=np.int32)] + data2 = [np.random.rand(16, 16, 16, 16).astype(dtype), np.array([1, 2], dtype=np.int32)] + + for data in [data0, data1, data2]: + testop(data) + testop(data, keep_dims=False) + testop(data, keep_dims=True) def _test_forward_reduce_quantized(testop):