[FRONTEND][ONNX] Error converting operator Transpose: TVMError: PermuteDims expects the number of input axes to equal the ndim of the input tensor.

[coffezhou]

Expected behavior

The onnx frontend should import the model correctly.

Actual behavior

Error converting operator Transpose, with inputs: [R.where(v0_0, v3_0, lv14)]
Traceback (most recent call last):
  File "/home/carla/Documents/test_tvm/test_tvm_using_onnx_model/random_test_tvm.py", line 172, in <module>
    main()
  File "/home/carla/Documents/test_tvm/test_tvm_using_onnx_model/random_test_tvm.py", line 161, in main
    check_correctness(onnx_model)
  File "/home/carla/Documents/test_tvm/test_tvm_using_onnx_model/random_test_tvm.py", line 104, in check_correctness
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3690, in from_onnx
    return g.from_onnx(graph, opset)
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3321, in from_onnx
    self._construct_nodes(graph)
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3501, in _construct_nodes
    raise err
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3498, in _construct_nodes
    op = self.bb.normalize(op)
  File "/home/carla/Documents/tvm/python/tvm/relax/block_builder.py", line 667, in normalize
    return _ffi_api.BlockBuilderNormalize(self, expr)  # type: ignore
  File "/home/carla/Documents/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 245, in __call__
    raise_last_ffi_error()
  File "/home/carla/Documents/tvm/python/tvm/_ffi/base.py", line 465, in raise_last_ffi_error
    raise py_err
tvm._ffi.base.TVMError: Traceback (most recent call last):
  8: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::RelaxExpr (tvm::relax::BlockBuilder, tvm::RelaxExpr const&)>::AssignTypedLambda<tvm::runtime::Registry::set_body_method<tvm::relax::BlockBuilder, tvm::relax::BlockBuilderNode, tvm::RelaxExpr, tvm::RelaxExpr const&, void>(tvm::RelaxExpr (tvm::relax::BlockBuilderNode::*)(tvm::RelaxExpr const&))::{lambda(tvm::relax::BlockBuilder, tvm::RelaxExpr const&)#1}>(tvm::runtime::Registry::set_body_method<tvm::relax::BlockBuilder, tvm::relax::BlockBuilderNode, tvm::RelaxExpr, tvm::RelaxExpr const&, void>(tvm::RelaxExpr (tvm::relax::BlockBuilderNode::*)(tvm::RelaxExpr const&))::{lambda(tvm::relax::BlockBuilder, tvm::RelaxExpr const&)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  7: tvm::relax::Normalizer::Normalize(tvm::RelaxExpr const&)
  6: tvm::relax::Normalizer::VisitExpr(tvm::RelaxExpr const&)
  5: _ZZN3tvm5relax11ExprFunctorIFNS_9RelaxExprERKS2_EE10InitVTableEvENUlRKNS_7r
  4: tvm::relax::Normalizer::VisitExpr_(tvm::relax::CallNode const*)
  3: tvm::relax::Normalizer::InferStructInfo(tvm::relax::Call const&)
  2: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::relax::StructInfo (tvm::relax::Call const&, tvm::relax::BlockBuilder const&)>::AssignTypedLambda<tvm::relax::StructInfo (*)(tvm::relax::Call const&, tvm::relax::BlockBuilder const&)>(tvm::relax::StructInfo (*)(tvm::relax::Call const&, tvm::relax::BlockBuilder const&))::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  1: tvm::relax::InferStructInfoPermuteDims(tvm::relax::Call const&, tvm::relax::BlockBuilder const&)
  0: tvm::relax::BlockBuilderImpl::ReportFatal(tvm::Diagnostic const&)
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 157
TVMError: PermuteDims expects the number of input axes to equal the ndim of the input tensor. However, the tensor ndim is 0 while the given number of axes is 4
[12:00:16] /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:64: Warning: BlockBuilder destroyed with remaining blocks!

Environment

OS: Ubuntu 20.04
TVM: 0.20.dev0 (6e8c367)

Steps to reproduce

This bug can be reproduced by the following code with the model in the attachment. For the model, it can be correctly checked by onnx.checker.check_model. However, the onnx frontend cannot import it.

import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import onnx

model_path = "model.onnx"
model = onnx.load(model_path)

onnx.checker.check_model(model, full_check=True)

tvm_model = from_onnx(model, keep_params_in_input=True)

model.zip

cc @KJlaccHoeUM9l