[Relay][Frontend][ONNX] Should flatten op is converted to batch_flatten ?

[Oldpan]

There is a bug in onnx frontend I think.
When I using onnx model which has a flatten op with axis attribute and the value is 1.
And then I run relay.frontend.from_onnx , it will occur a issue :

batch_flatten() got an unexpected keyword argument 'axis'

So I modify the source code in .../relay/op/nn/nn.py - line 429.
Change

def batch_flatten(data):

to

def batch_flatten(data, axis=1):

And it works but obviously this is a temporary solution.
I see this pull request #2843 had renamed 'Flatten' to 'batch_flatten'.
Although flatten with default axis 1 is just same as batch_flatten but I think this behavior should be regulated. Probably we can use reshape op to replace flatten and don't rename flatten :

def flatten(data, axis=1):

    newshape = [0]*(axis+1)
    newshape[axis] = -1;
    return _make.reshape(data, list(newshape))

It can work but I'm not sure it is suitable.

[tqchen]

Thanks for bringing this up, we can do an optional conversion(convert to batch_flatten when it matches, use reshape otherwise) We can also enhance batch_flatten. Contribution is welcomed!

[Oldpan]

Thanks for reply.
I'LL try a contribution for this

[tqchen]

#3180 Thanks @Oldpan