[Relay][ONNX][Fix] Flatten in OnnxConverter

[ganler]

Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.

Simple Flatten operator's shape cannot be evaluated when using ONNX. This is because the flatten op does not try to do constant folding thus creating unnecessary dynamic operator, causing graph executor unable to execute it.

To trigger this bug, simply use this script to generate a ONNX model and run it with graph executor.

import torch

N = 1 # can be something else

class Net(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
    
    def forward(self, x):
        return x.flatten()

net = Net().eval()

o = net(torch.zeros((N), dtype=torch.float))
print(o.shape)

with torch.no_grad():
  torch.onnx.export(net, (
      torch.ones((N), dtype=torch.float)), "output.onnx", verbose=True, opset_version=14)

It will throw errors: the relay script is unnecessarily complicated.

def @main(%v0: Tensor[(1), float32]) -> Tensor[(?, ?), float32] {
  %0 = shape_of(%v0, dtype="int32") /* ty=Tensor[(1), int32] */;
  %1 = strided_slice(%0, begin=[0], end=[0], strides=[1], axes=None) /* ty=Tensor[(0), int32] */;
  %2 = strided_slice(%0, begin=[0], end=[1], strides=[1], axes=None) /* ty=Tensor[(1), int32] */;
  %3 = prod(%1, keepdims=True) /* ty=Tensor[(1), int32] */;
  %4 = prod(%2, keepdims=True) /* ty=Tensor[(1), int32] */;
  %5 = (%3, %4);
  %6 = concatenate(%5) /* ty=Tensor[(2), int32] */;
  dyn.reshape(%v0, %6, newshape=[]) /* ty=Tensor[(?, ?), float32] */
}

Traceback (most recent call last):
...
    executor = relay.build_module.create_executor(
  File "/home/ganler/Documents/tvm/python/tvm/relay/backend/interpreter.py", line 171, in evaluate
    return self._make_executor()
  File "/home/ganler/Documents/tvm/python/tvm/relay/build_module.py", line 621, in _make_executor
    raise ValueError(
ValueError: ('Graph Executor only supports static graphs, got output type', TensorType([?, ?], float32))

It looks much better after this fix.

def @main(%v0: Tensor[(1), float32]) -> Tensor[(1, 1), float32] {
  reshape(%v0, newshape=[1, 1]) /* ty=Tensor[(1, 1), float32] */
}

cc: @AndrewZhaoLuo @vinx13 @jwfromm @masahi

[masahi]

very nice!

[ganler]

It seems the CI is broken. I will push an empty commit when it's fixed.

[masahi]

Hopefully fixed by #10610

[masahi]

The CI issue has been fixed thanks to #10626

[ganler]

@masahi Thank you! The CI now passes!