[ONNX] Fix Split operator for uneven splits with num_outputs

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

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]: