[TFLite] added scalar axis value handling in reduce

Axis value in reduce can now be specified as scalar

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

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):