[Torch] Add support for split #5174
Conversation
| indices.append(now_indice) | ||
| now_indice += split_size | ||
|
|
||
| return _op.split(data, indices, dim) |
There was a problem hiding this comment.
Because this op is splitting the index evenly, you can just provide a number of sections to split to rather than a list of indices.
return _op.split(data, int(_infer_shape(data)[dim] / split_size), dim)
There was a problem hiding this comment.
The split op in Relay and PyTorch have different behaviors. In Relay, if indices_or_sections is an integer, the input will be divided equally along the given axis. But in Pytorch, if split_size_or_sections is an integer type, then tensor will be split into equally sized chunks (if possible). The last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.
python/tvm/relay/frontend/pytorch.py
Outdated
| data = inputs[0] | ||
| dim = int(inputs[2]) | ||
|
|
||
| now_indice = 0 |
There was a problem hiding this comment.
A better name for this variable would be split_index.
|
|
||
| class Split2(Module): | ||
| def forward(self, *args): | ||
| return torch.split(args[0], [2, 3, 5], 1) |
There was a problem hiding this comment.
Does this end up testing split_with_sizes?
There was a problem hiding this comment.
Yes, the Split2 will test split_with_sizes.
|
Thanks @wyc-ruiker |
* [Torch] Add support for split * fix * fix test class
* [Torch] Add support for split * fix * fix test class
* [Torch] Add support for split * fix * fix test class
Thanks for contributing to TVM! Please refer to guideline https://docs.tvm.ai/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.
add split in #5133 @masahi @alexwong @jwfromm