From 836a80adabf6d51eb9906e5e192252db4d67fae5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 18 Mar 2026 18:05:30 +0000 Subject: [PATCH 1/2] Initial plan From 0a4f8bbb75c376039fdee3b0924f19a19365b82b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 18 Mar 2026 18:13:10 +0000 Subject: [PATCH 2/2] Fix SplitToSequence with unequal scalar split generating incorrect Split node When split is a scalar that doesn't evenly divide the axis dimension, create an explicit split sizes tensor and pass it to Split, instead of using num_outputs alone (which would produce an equal split). Example: input [1,8400,80] with scalar split=5000 on axis=1 now correctly produces [1,5000,80] and [1,3400,80] instead of [1,4200,80] and [1,4200,80]. Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> --- onnxscript/optimizer/_constant_folding.py | 27 +++++++++++-- .../optimizer/_constant_folding_test.py | 39 +++++++++++++++++++ 2 files changed, 62 insertions(+), 4 deletions(-) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 574ddd8aef..ed69487c6c 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -860,11 +860,30 @@ def split_to_sequence(node: ir.Node, op, state: OptimizerState) -> ReturnValue: split_dimension_size = shape[axis] if not isinstance(split_dimension_size, int): return None - num_outputs = math.ceil(split_dimension_size / split_value.item()) + split_size = int(split_value.item()) + num_outputs = math.ceil(split_dimension_size / split_size) split_outputs = [f"{output.name}_split_{i}" for i in range(num_outputs)] - split_values = op.Split( - input, axis=axis, num_outputs=num_outputs, _outputs=split_outputs - ) + if split_dimension_size % split_size != 0: + # Uneven split: the last chunk is smaller. We must pass explicit split + # sizes to Split, because Split with only num_outputs would do an + # equal (or near-equal) split ignoring the original chunk size. + remainder = split_dimension_size - (num_outputs - 1) * split_size + explicit_split_sizes = [split_size] * (num_outputs - 1) + [remainder] + explicit_split = op.Constant( + value_ints=explicit_split_sizes, + _outputs=[f"{output.name}_split_sizes"], + ) + split_values = op.Split( + input, + explicit_split, + axis=axis, + num_outputs=num_outputs, + _outputs=split_outputs, + ) + else: + split_values = op.Split( + input, axis=axis, num_outputs=num_outputs, _outputs=split_outputs + ) else: return None diff --git a/onnxscript/optimizer/_constant_folding_test.py b/onnxscript/optimizer/_constant_folding_test.py index 080af9c2f3..2f68aabba3 100644 --- a/onnxscript/optimizer/_constant_folding_test.py +++ b/onnxscript/optimizer/_constant_folding_test.py @@ -284,6 +284,45 @@ def test_static_split_to_sequence_with_scalar_split_and_squence_at_is_folded_as_ self.assertEqual(len(optimized.graph[-2].outputs), 4) self.assertEqual(optimized.graph[-2].op_type, "Split") + def test_static_split_to_sequence_with_unequal_scalar_split_and_sequence_at_is_folded_as_split( + self, + ): + """Test that an unequal scalar split is preserved correctly (not turned into equal split). + + Regression test for: SplitToSequence with scalar split that doesn't evenly divide + the axis dimension should produce a Split with explicit split sizes, not an equal split. + E.g., splitting dim=8400 with split=5000 should produce [5000, 3400], not [4200, 4200]. + """ + model = """ +< + ir_version: 8, + opset_import: ["" : 18] +> +func (float[1,8400,80] x) => (float[1,N,80] return_val) { + int64_5000 = Constant () + splits = SplitToSequence (x, int64_5000) + int64_0 = Constant () + split_0 = SequenceAt (splits, int64_0) + int64_1 = Constant () + split_1 = SequenceAt (splits, int64_1) + return_val = Concat (split_0, split_1) +}""" + + optimized = self._fold(model) + split_nodes = [n for n in optimized.graph if n.op_type == "Split"] + self.assertEqual(len(split_nodes), 1) + split_node = split_nodes[0] + self.assertEqual(len(split_node.outputs), 2) + # The Split node must have an explicit split input (not just num_outputs), + # so that the split is [5000, 3400] and not [4200, 4200]. + split_sizes_input = split_node.inputs[1] + self.assertIsNotNone(split_sizes_input, "Split node must have explicit split sizes") + # Verify the actual split sizes are [5000, 3400], not [4200, 4200] + self.assertIsNotNone(split_sizes_input.const_value) + np.testing.assert_array_equal(split_sizes_input.const_value.numpy(), [5000, 3400]) + # Check no SequenceAt remains + self.assertTrue(all(n.op_type != "SequenceAt" for n in optimized.graph)) + def test_static_split_to_sequence_with_list_split_and_squence_at_is_folded_as_split( self, ):