Repository: apache/tvm

Issue:
[Bug] Softmax in TFLite converter is channel first instead of channel last
Compiled a TFLite toy model (conv + softmax) using TVM0.7 to generate a softmax output.

### Expected behavior
The uint8 output distribution of the output should look something like (below is tflite runtime output)

![tflite_hist](https://user-images.githubusercontent.com/7594901/134404281-3752887b-da54-4d58-b461-1f867a6a3f4f.png)

### Actual behavior
![tvm_hist](https://user-images.githubusercontent.com/7594901/134404357-94463c5d-0059-47c4-a790-eafccbf01b41.png)

After I change axis to 3 [here](https://github.com/apache/tvm/blob/701d2c32759c95951c4eeca229addd1036539698/python/tvm/relay/frontend/tflite.py#L684), I get output that is similar to the expected output. 

### Environment
OS: Ubuntu 20.04
TVM: 0.7

### Steps to reproduce
deps: tvm 0.7, numpy, and matplotlib to generate the plot
The zip contains the tflite model. Feel free to use a random 

[tflite_model.tflite.zip](https://github.com/apache/tvm/files/7213238/tflite_model.tflite.zip)
[Here](https://drive.google.com/file/d/1FJLyOVhFbLNWElnqNu9Ur-ki3WEXWDdF/view?usp=sharing) is the input numpy array
```
import tvm
import numpy as np
import tflite
import tflite_runtime.interpreter as tflite_interpreter

from tvm import relay, transform
from tvm.contrib import graph_runtime as runtime

import matplotlib.pyplot as plt


def generate_hist(x, title, filename):
    plt.title(title)
    plt.xlabel('uint8 value')
    plt.xlim([-1,256])
    plt.ylabel('frequency')
    H, bins = np.histogram(x.flatten(), bins=256, range=(0,256))
    plt.bar(bins[:-1], H)
    plt.savefig(filename)
    plt.clf()

# TFLite compilation
model_path = 'tflite_model.tflite'
interpreter = tflite_interpreter.Interpreter(model_path=model_path)

interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

with open('conv_softmax_input.bin', 'rb') as f:
    input_arr = np.load(f) * 255
input_data = input_arr.astype(np.uint8)

interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

tflite_output1 = interpreter.get_tensor(output_details[0]['index'])


# TVM Compilation
tflite_model_buf = open(model_path, "rb").read()
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)

input_dict={"input":(1,100,160,256)}
input_dtype = {"input":"uint8"}
mod, params = relay.frontend.from_tflite(
    tflite_model, shape_dict=input_dict, dtype_dict=input_dtype
)

target = "llvm"
with transform.PassContext(opt_level=2):
    lib = relay.build(mod, target, params=params)

#invokes pure TVM compilation
module = runtime.GraphModule(lib["default"](tvm.cpu()))
module.set_input("input", tvm.nd.array(input_data))
module.run()
tvm_output1 = module.get_output(0).asnumpy()

generate_hist(tflite_output1, 'tflite quant(uint8) hist', 'tflite_hist.png')
generate_hist(tvm_output1, 'tvm quant(uint8) hist', 'tvm_hist.png')
```


Discussion:

Resolving patch:
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -1496,7 +1496,7 @@ def _impl_v13(cls, inputs, attr, params):
             axis = relay.TupleGetItem(axes, i)
             # Unpack scalar
             axis = relay.reshape(axis, [])
-            axis = relay.If(
+            axis = relay.where(
                 axis >= relay.const(0, "int64"), axis, axis + relay.const(rank_input, "int64")
             )
             result = _op.expand_dims(result, axis)
@@ -1509,12 +1509,18 @@ class Squeeze(OnnxOpConverter):
     @classmethod
     def _impl_v1(cls, inputs, attr, params):
         axis = attr.get("axes", None)
-        return _op.squeeze(*inputs, axis)
+        return _op.squeeze(inputs[0], axis)
 
     @classmethod
     def _impl_v13(cls, inputs, attr, params):
         axis = inputs[1]
         dtype = infer_type(axis).checked_type.dtype
+
+        if isinstance(axis, _expr.Constant):
+            constant_axes = list(inputs[1].data.numpy())
+            constant_axes = list(map(int, constant_axes))
+            return _op.squeeze(inputs[0], constant_axes)
+
         rank = _op.shape_of(_op.shape_of(inputs[0], dtype), dtype)
         axis = _op.where(axis < _op.const(0, dtype), axis + rank, axis)
         return _op.squeeze(inputs[0], fold_constant(axis))
@@ -1640,7 +1646,7 @@ def normalize_gather_indices(data, indices, axis):
     """Make sure gather indicies aren't negative"""
     ind_dtype = infer_type(indices).checked_type.dtype
     # Normalize the indices to a positive range
-    s = _op.take(_op.shape_of(data, dtype=ind_dtype), _op.const(axis))
+    s = _op.take(_op.shape_of(data, dtype=ind_dtype), _op.const(axis, dtype="int64"))
     cond = fold_constant(indices < _op.const(0, ind_dtype))
     if isinstance(cond, _expr.Constant):
         val = cond.data.numpy()


Test patch:
diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -4939,25 +4939,8 @@ def verify_eyelike(indata):
     "test_maxpool_with_argmax_2d_precomputed_strides",
     "test_maxunpool_export_with_output_shape",
     "test_mvn",
-    # When unsqueeze is fully supported, remaining nllloss tests should work:
-    "test_nllloss_NC_expanded",
-    "test_nllloss_NCd1_expanded",
-    "test_nllloss_NCd1_ii_expanded",
-    "test_nllloss_NCd1_mean_weight_negative_ii_expanded",
-    "test_nllloss_NCd1_weight_expanded",
-    "test_nllloss_NCd1_weight_ii_expanded",
-    "test_nllloss_NCd1d2_expanded",
-    "test_nllloss_NCd1d2_no_weight_reduction_mean_ii_expanded",
-    "test_nllloss_NCd1d2_reduction_mean_expanded",
-    "test_nllloss_NCd1d2_reduction_sum_expanded",
-    "test_nllloss_NCd1d2_with_weight_expanded",
-    "test_nllloss_NCd1d2_with_weight_reduction_mean_expanded",
-    "test_nllloss_NCd1d2_with_weight_reduction_sum_expanded",
-    "test_nllloss_NCd1d2_with_weight_reduction_sum_ii_expanded",
+    # This test fails llvm with a lowering error:
     "test_nllloss_NCd1d2d3_none_no_weight_negative_ii_expanded",
-    "test_nllloss_NCd1d2d3_sum_weight_high_ii_expanded",
-    "test_nllloss_NCd1d2d3d4d5_mean_weight_expanded",
-    "test_nllloss_NCd1d2d3d4d5_none_no_weight_expanded",
     "test_qlinearmatmul_2D",
     "test_qlinearmatmul_3D",
     "test_range_float_type_positive_delta_expanded",