Split-18 performs differently compared to np.array_split, torch.tensor_split for last element in uneven split

[take-cheeze]

Bug Report

Is the issue related to model conversion?

No

Describe the bug

In uneven Split-18, only the last element would be less than others.
Though in np.array_split/torch.tensor_split it would make target axis length of each element at least 1 less.
With this difference, exporting split operation to opset 18 would be a hard job

System information

• ONNX version : 1.13.0

Reproduction instructions

tensor_split / array_split runs like:

>>> torch.arange(10).tensor_split(4)
(tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7]), tensor([8, 9]))

But Split-18 will split it like:

(tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7, 8]), tensor([9]))

Also reference implementation should act same:

onnx/onnx/reference/ops/op_split.py

Line 22 in 9125ac3

split[-1] += mat.shape[axis] - sum(split) # type: ignore

Expected behavior

There should be a mode attribute

Notes

[xadupre]

The current implementation follows the specifications.

Split a tensor into a list of tensors, along the specified ‘axis’. Either input ‘split’ or the attribute ‘num_outputs’ should be specified, but not both. If the attribute ‘num_outputs’ is specified, then the tensor is split into equal sized parts. If the tensor is not evenly splittable into num_outputs, the last chunk will be smaller. If the input ‘split’ is specified, it indicates the sizes of each output in the split.

You suggest adding a new attribute mode to follow torch's behaviour?

PR #4743 fixes the implementation but do not change the behaviour.

[NValerij]

It is also unclear from specification what should happen in this particular case: split 9 elements into 4 outputs: [3, 3, 3, 0]?.
Is it a error if empty tensor occurs in output?

[xadupre]

As weird as it seems, that's what the specifications says.

[gramalingam]

Agree that it would be useful to add a flag to get the torch behavior.

[gramalingam]

What about SplitToSequence and its usage in the Torch exporter? Does it have a similar requirement? @justinchuby ?

[justinchuby]

We tested with torch.split and it seems fine. torch.tensor_split seems to behave differently

[skottmckay]

The Split-18 spec is broken for some combinations.

e.g. split 5 into 4 is impossible as you'd need 2, 1, 1, 1 as 2, 2, 2, 0 has too many elements.

[zhenhuaw-me]

I think @liqunfu 's proposal in #5766 (comment) makes more sense, as a mode such as mode:torch requires background of PyTorch API definition.

One step further, we can fix the current attribute num_outputs by aligning it with numpy.array_split and torch.tensor_split

If indices_or_sections is an integer n or a zero dimensional long tensor with value n, input is split into n sections along dimension dim. If input is divisible by n along dimension dim, each section will be of equal size, input.size(dim) / n. If input is not divisible by n, the sizes of the first int(input.size(dim) % n) sections will have size int(input.size(dim) / n) + 1, and the rest will have size int(input.size(dim) / n).