[onnx] Fix onnx-to-torch lowering for flatten shape

[rsuderman]

The existing flatten lowering did not define what the intermediate
shape was. This could result in failures to lower further to linalg as
the intermediate shape was unknown. Added a shape refinement section.

[newling]

Did this come up in a downstream test? Would this be caught with #2795 landed?

[rsuderman]

Did this come up in a downstream test? Would this be caught with #2795 landed?

#2795 will help catch these tests in the future. It was caught by attempting to push the first model through (resnet) and we found this error, along with a few others.

[newling]

Just the 2 issues

1. axis == rank
2. handling dynamic

[rsuderman]

Just the 2 issues

1. axis == rank
2. handling dynamic

If axis == rank then the loop never executes meaning we never index into shape.

[newling]

The bool dynamic = . line needs another look

[renxida]

It looks like at DefaultDomainAtoF.cpp:1424 there's a redundant or. Did you mean shape[i]?

Code:

         bool dynamic = shape[axis] == Torch::kUnknownSize ||
                         shape[axis] == Torch::kUnknownSize;

otherwise lgtm

[rsuderman]

The bool dynamic = . line needs another look

Woops, fixed it up. Just needed to check whether the value is dynamic.

[newling]

looks all good to me now