[Bug][FRONTEND][ONNX] Error converting operator Expand: TVMError: broadcast_to expects the input tensor shape is broadcastable to the target shape.

[coffezhou]

Expected behavior

The onnx frontend should import the model correctly.

Actual behavior

Error converting operator Expand, with inputs: [R.mean(lv4, axis=[2], keepdims=False), metadata["relax.expr.Constant"][0]
# Metadata omitted. Use show_meta=True in script() method to show it.]
Traceback (most recent call last):
  File "/home/carla/Documents/test_tvm/0318/test_relax1.py", line 14, in <module>
    tvm_model = from_onnx(model, 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 "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 270, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 259, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
  File "/home/carla/Documents/tvm/python/tvm/_ffi/base.py", line 468, in raise_last_ffi_error
    raise py_err
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 533, in tvm::relax::Normalizer::Normalize(tvm::RelaxExpr const&)
    Expr normalized = this->VisitExpr(expr);
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 616, in tvm::relax::Normalizer::VisitExpr(tvm::RelaxExpr const&)
    return ExprFunctor::VisitExpr(expr);
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 664, in tvm::relax::Normalizer::VisitExpr_(tvm::relax::CallNode const*)
    auto inferred_sinfo = InferStructInfo(call);
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 847, in tvm::relax::Normalizer::InferStructInfo(tvm::relax::Call const&)
    return op_map_infer_struct_info_[op](call, GetRef<BlockBuilder>(this));
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/src/relax/op/tensor/manipulate.cc", line 95, in tvm::relax::InferStructInfoBroadcastTo(tvm::relax::Call const&, tvm::relax::BlockBuilder const&)
    ctx->ReportFatal(
                ^^^^^^
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 157, in tvm::relax::BlockBuilderImpl::ReportFatal(tvm::Diagnostic const&)
    LOG(FATAL) << diagnostic->message;
                  ^^^^^^^^^^^^^^^^^^^^^
tvm._ffi.base.TVMError: Traceback (most recent call last):
  6: tvm::relax::Normalizer::Normalize(tvm::RelaxExpr const&)
        at /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:533
  5: tvm::relax::Normalizer::VisitExpr(tvm::RelaxExpr const&)
        at /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:616
  4: non-virtual thunk to tvm::relax::Normalizer::VisitExpr_(tvm::relax::CallNode const*)
  3: tvm::relax::Normalizer::VisitExpr_(tvm::relax::CallNode const*)
        at /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:664
  2: tvm::relax::Normalizer::InferStructInfo(tvm::relax::Call const&)
        at /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:847
  1: tvm::relax::InferStructInfoBroadcastTo(tvm::relax::Call const&, tvm::relax::BlockBuilder const&)
        at /home/carla/Documents/tvm/src/relax/op/tensor/manipulate.cc:95
  0: tvm::relax::BlockBuilderImpl::ReportFatal(tvm::Diagnostic const&)
        at /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:157
  File "/home/carla/Documents/tvm/src/relax/ir/block_builder.cc", line 157
TVMError: broadcast_to expects the input tensor shape is broadcastable to the target shape. The target shape at dim 3 is T.int64(56) while the input tensor shape at dim 1 is T.int64(25), which are not equal.
[14:42:56] /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 (f6236ce)

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