Operator spec of Split operator's attribute num_outputs is wrong

[zhenhuaw-me]

The Problem

In op version 18, we added the spec of Split attribute num_outputs with the statement

num_outputs - INT : Number of outputs to split parts of the tensor into. If the tensor is not evenly splittable the last chunk will be smaller.

Consider the case Split(input, 4) where input has dim [5], what the shape of the outputs will be?

According to the shape inference code, the shape of the output tensors should be [2], [2], [2], [-1] which is illegal.

                int chunk_size = (split_dim_value / num_outputs) + 1;
                int last_chunk_size = split_dim_value - (chunk_size * (num_outputs - 1));
                split.resize(num_outputs - 1, chunk_size);
                split.push_back(last_chunk_size);

Reference

It seems to me that we were trying to combine the definition of numpy/TF and Torch.

For Numpy/TF, they use the similar semantic, but REJECT if not divisible:

If num_or_size_splits is an int, then it splits value along the dimension axis into num_or_size_splits smaller tensors. This requires that value.shape[axis] is divisible by num_or_size_splits.

Torch instead declares it as a split size, in which way indivisible is acceptable:

If split_size_or_sections is an integer type, then tensor will be split into equally sized chunks (if possible). Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

Possible Solutions

1. Keep the num_outputs and reject if indivisible.
2. Change to Torch's definition, thus we need to change the attribute name.

[zhenhuaw-me]

The Split-18 PR: #4481 @p-wysocki

What's the intention of num_outputs, as it is different from both Torch and TensorFlow.

cc @onnx/sig-operators-approvers

[liqunfu]

Good catch! The onnx spec need to be updated to require the input being evenly splittable by num_outputs, if num_outputs is specified. This is the same as tf: "that value.shape[axis] is divisible by num_or_size_splits". In the case where num_or_size_splits in TF and split_size_or_sections in pytorch NOT being a single integer, it is covered by "split" input in ONNX. To support split_size_or_sections being a single integer in Pytorch, we need add an attribute "split_size" where the last output may be smaller than split_size.

Here is how cases in TF and Pytroch are covered in ONNX:
Pytorch:
input = [1, 2, 3, 4, 5], split_size_or_sections = 2 => output = ([1, 2], [3, 4], [5]); // ONNX split_size = 2
input = [1, 2, 3, 4, 5], split_size_or_sections = [2, 3] => output = ([1, 2], [3, 4, 5]) // ONNX split = [2, 3]
TF:
input = [1, 2, 3, 4, 5], num_or_size_splits = 2 => no dividable error; // ONNX no dividable error
input = [1, 2, 3, 4, 5, 6], num_or_size_splits = 2 => output = ([0, 1, 2], [3, 4, 5]; // ONNX num_outputs = 2
input = [1, 2, 3, 4, 5], num_or_size_splits = [2, 3] => output = ([1, 2], [3, 4, 5]) // ONNX split = [2, 3]

[zhenhuaw-me]

@liqunfu Thank you for the update! That sounds good. Do we plan to include this in next release?

[liqunfu]

@liqunfu Thank you for the update! That sounds good. Do we plan to include this in next release?

Yes. It however needs to consolidate with #4742 the PR: #5321.

@p-wysocki , any update of the PR and shall/does it cover this issue?

[zhenhuaw-me]

Copy #4742 (comment)

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).