From d909bf07bf61c39feed349fbb687260b84175808 Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com> Date: Fri, 13 Feb 2026 15:09:59 +0530 Subject: [PATCH 1/2] [ONNX] Fix Split operator for uneven splits with num_outputs The ONNX frontend's Split converter failed when the tensor dimension wasn't evenly divisible by num_outputs (e.g., splitting a length-10 tensor into 3 parts). This implements the ceil-based split calculation from the ONNX Opset 18 spec: block_size = ceil(dim / num_outputs), yielding [block_size] * (N-1) + [remainder] for the output sizes. When the input shape is statically known, explicit cumulative split indices are computed and passed to relax.op.split. For dynamic shapes, the integer num_outputs is passed through as before. Fixes #18751 Signed-off-by: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com> --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 37 ++++++++++++++++++- tests/python/relax/test_frontend_onnx.py | 16 +++++++- 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 6d38d2b2ca8c..42f6e3ebea83 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -2531,6 +2531,35 @@ def _impl_v1(cls, bb, inputs, attr, params): class Split(OnnxOpConverter): """Converts an onnx Split node into an equivalent Relax expression.""" + @classmethod + def _compute_split_indices(cls, data, axis, num_outputs): + """Compute split indices for potentially uneven splits. + + Per ONNX Opset 18 spec, when num_outputs is given without explicit + split sizes, the split uses ceil-based block sizes: + block_size = ceil(dim / num_outputs) + The first (N-1) chunks have size block_size, and the last chunk + gets the remainder. + + ``relax.op.split`` with an integer divides the axis evenly, which is + only correct when the dimension is divisible by ``num_outputs``. For an + uneven split on a statically known axis this returns explicit cumulative + indices instead. + + The integer form is kept whenever it is already correct -- an even split, + or a dynamic axis where the sizes are not known at compile time -- so the + emitted IR is unchanged for every case that worked before. + """ + shape = data.struct_info.shape + if shape is not None and isinstance(shape, relax.ShapeExpr): + dim_val = shape[axis] + if isinstance(dim_val, (tir.IntImm,)): + dim = int(dim_val) + if dim % num_outputs: + block_size = math.ceil(dim / num_outputs) + return [block_size * i for i in range(1, num_outputs)] + return num_outputs + @classmethod def _impl_v1(cls, bb, inputs, attr, params): splits = attr.get("split", None) @@ -2542,7 +2571,9 @@ def _impl_v1(cls, bb, inputs, attr, params): indices.append(index) # When splits isnt specified divide evenly over axis. else: - indices = attr["tvm_custom"]["num_outputs"] + axis = attr.get("axis", 0) + num_outputs = attr["tvm_custom"]["num_outputs"] + indices = cls._compute_split_indices(inputs[0], axis, num_outputs) return relax.op.split(inputs[0], indices, attr.get("axis", 0)) @classmethod @@ -2563,7 +2594,9 @@ def _impl_v13(cls, bb, inputs, attr, params): raise ValueError("Dynamic Split not yet supported") # When splits isnt specified divide evenly over axis. else: - indices = attr["tvm_custom"]["num_outputs"] + axis = attr.get("axis", 0) + num_outputs = attr.get("num_outputs", attr["tvm_custom"]["num_outputs"]) + indices = cls._compute_split_indices(inputs[0], axis, num_outputs) return relax.op.split(inputs[0], indices, attr.get("axis", 0)) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 730a969b626b..59e16bbe44b5 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -7382,7 +7382,16 @@ def expected_input_shape(shape): dtype = np.dtype(fp_arith).name input_shape = expected_input_shape(indata_shape) if not pass_split: - indices_or_sections = len(outdata_shapes) + n_outputs = len(outdata_shapes) + axis_dim = shape_tuple(indata_shape)[axis] + # relax.op.split with an integer splits evenly, so an uneven + # num_outputs split lowers to explicit indices instead. Only when + # the axis is static: a dynamic one keeps the integer form. + if not dynamic and axis_dim % n_outputs: + block_size = -(-axis_dim // n_outputs) + indices_or_sections = [block_size * i for i in range(1, n_outputs)] + else: + indices_or_sections = n_outputs elif len(outdata_shapes) == 1: indices_or_sections = 1 else: @@ -7448,6 +7457,11 @@ def main(input: R.Tensor(input_shape, dtype=dtype)): (1, [[1]], [1], 0, True, 11), ((1, 2), [[2]], [2], 1, True, 11), ((1, 2), [[2]], [1], 0, True, 11), + # Uneven num_outputs splits (opset 18): ceil-based blocks per the spec. + # 10 / 3 -> [4, 4, 2] + (10, [[4], [4], [2]], False, 0, False, 18), + # 7 / 3 along axis 1 -> [3, 3, 1] + ((4, 7), [[4, 3], [4, 3], [4, 1]], False, 1, False, 18), ] for fp_arith in [np.float16, np.float32]: From 4903e3c026c5b88ded5e4f260101b961e24303d7 Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com> Date: Tue, 2 Jun 2026 23:41:45 +0530 Subject: [PATCH 2/2] refactor(onnx): reuse `axis` local in Split._impl_v1/v13 Per @gemini-code-assist's review: avoid the redundant `attr.get("axis", 0)` lookup at the split call site by reusing the `axis` local already retrieved for the uneven-split branch. Pure refactor, no behavior change. The CPU build error mentioned by @tlopex is on the C++ side and unrelated to this change. --- python/tvm/relax/frontend/onnx/onnx_frontend.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 42f6e3ebea83..41fbe76487a0 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -2562,6 +2562,7 @@ def _compute_split_indices(cls, data, axis, num_outputs): @classmethod def _impl_v1(cls, bb, inputs, attr, params): + axis = attr.get("axis", 0) splits = attr.get("split", None) if splits is not None and len(splits) > 1: indices = [] @@ -2571,13 +2572,13 @@ def _impl_v1(cls, bb, inputs, attr, params): indices.append(index) # When splits isnt specified divide evenly over axis. else: - axis = attr.get("axis", 0) num_outputs = attr["tvm_custom"]["num_outputs"] indices = cls._compute_split_indices(inputs[0], axis, num_outputs) - return relax.op.split(inputs[0], indices, attr.get("axis", 0)) + return relax.op.split(inputs[0], indices, axis) @classmethod def _impl_v13(cls, bb, inputs, attr, params): + axis = attr.get("axis", 0) splits = inputs[1] splits_rank = None if splits is not None: @@ -2594,10 +2595,9 @@ def _impl_v13(cls, bb, inputs, attr, params): raise ValueError("Dynamic Split not yet supported") # When splits isnt specified divide evenly over axis. else: - axis = attr.get("axis", 0) num_outputs = attr.get("num_outputs", attr["tvm_custom"]["num_outputs"]) indices = cls._compute_split_indices(inputs[0], axis, num_outputs) - return relax.op.split(inputs[0], indices, attr.get("axis", 0)) + return relax.op.split(inputs[0], indices, axis) def get_prim_value_list(values):