[Relay][Pytorch] Add support for aten::unflatten

[mshr-h]

Fix #15663
Support torch.unflatten.

cc @jikechao @vvchernov @Hzfengsy @junrushao

[vvchernov]

Hello @mshr-h! LGTM, but it is needed one check

[vvchernov]

Add check that dshape[dim] == multiplication of dimensions in unflattened_size

[mshr-h]

@vvchernov Thanks!
I don't think we have to check it because torch.jit.trace does it.
They provide something like the below RuntimeError when the shape is wrong.

RuntimeError: unflatten: Provided sizes [3, 5, 3, -1] don't multiply up to the size of dim 0 (60) in the input tensor

Should we add the check in TVM's PyTorch frontend?

[vvchernov]

Hello @mshr-h! Ok torch.jit.trace do it, but in this case we do not need assert len(dshape) > dim.
It looks like TVM usually rechecks all corner cases.
About -1 in unflattened_size. In this case we can multiply together other dimension and check dshape[dim] % mult == 0

[mshr-h]

@vvchernov Thanks. Added the assertion.

[vvchernov]

Hello @mshr-h! There are two cases: unflattened_size has -1 and does not have. You check only the first one. Example: dshape[dim] = 8, unflattened_size = [2, 2, 1, 1] pass your assert, but it is failure case

[vvchernov]

I've checked that _op.reshape does not check size correctness by it-self. Nevertheless I found that dim can be not only -1, but from list {0, -1, -2, -3, -4} See

tvm/python/tvm/relay/op/transform.py

Line 243 in 748882a

def reshape(data, newshape, allowzero=False):

I suggest to check -1 case only. It looks like torch does not have other options

[mshr-h]

@vvchernov
Thanks! Updated the shape check.

[mshr-h]

Ah, dim can be negative but I didn't check that. I'll add the check.

[mshr-h]

Ah, dim can be negative but I didn't check that. I'll add the check.

Done.

[jikechao]

LGTM! @mshr-h Thanks for your PR.

[vvchernov]

Hello @mshr-h! Thank you for careful work! Sorry for many small boring corrections. LGTM