diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 6d38d2b2ca8c..41fbe76487a0 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -2531,8 +2531,38 @@ 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): + axis = attr.get("axis", 0) splits = attr.get("split", None) if splits is not None and len(splits) > 1: indices = [] @@ -2542,11 +2572,13 @@ 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"] - return relax.op.split(inputs[0], indices, 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, 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: @@ -2563,8 +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: - indices = attr["tvm_custom"]["num_outputs"] - return relax.op.split(inputs[0], indices, 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, axis) def get_prim_value_list(values): 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]: