From fabccf9536a6add0744872dda403433b9dcc41bd Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Sun, 1 Nov 2026 13:49:36 +0530 Subject: [PATCH 1/6] test(rewriter): add regression tests for slice axis bounds in collapse_slices --- .../rules/common/_collapse_slices_test.py | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices_test.py b/onnxscript/rewriter/rules/common/_collapse_slices_test.py index 2c8101a875..06d11654d7 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices_test.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices_test.py @@ -178,3 +178,81 @@ def test_multi_element_steps_numerical_correctness(self): model, (np.random.rand(4, 5).astype(np.float32),), ) + + def test_slice_out_of_bounds_axis_does_not_crash_or_collapse(self): + """Slice with out-of-bounds axis should be rejected without IndexError.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_negative_out_of_bounds_axis_does_not_crash_or_collapse(self): + """Slice with negative out-of-bounds axis should be rejected without IndexError.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_out_of_bounds_axis_with_int64_max_ends_does_not_collapse(self): + """Slice with out-of-bounds axis must not collapse even if ends reaches int64_max.""" + model = ir.from_onnx_text( + f""" + + agraph (float[4, 8] data) => (float[4, 8] output) + {{ + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + }} + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_valid_negative_axis_collapses(self): + """Slice with a valid negative axis (e.g. -1 for rank 2) should collapse.""" + model_text = """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertIn("Identity", [node.op_type for node in model.graph]) + testing.assert_numerically_equal( + original, + model, + (np.random.rand(4, 8).astype(np.float32),), + ) From 27f1ff6d2eff9b218765a6246926c23930b0ec21 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Sun, 1 Nov 2026 13:51:09 +0530 Subject: [PATCH 2/6] fix(rewriter): validate and normalize axis in collapse_slices --- .../rewriter/rules/common/_collapse_slices.py | 39 +++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices.py b/onnxscript/rewriter/rules/common/_collapse_slices.py index 712aa410f3..2a2d36b313 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices.py @@ -55,13 +55,24 @@ def _check_if_redundant_slice( if starts_const.numpy().item() != 0: logger.info("The value 'start' is not 0.") return False + + axis = int(axes_const.numpy().item()) + if data.shape is not None: + rank = data.shape.rank() + if rank is not None: + if axis < -rank or axis >= rank: + logger.info("The value 'axis' is out of range for the data rank.") + return False + if axis < 0: + axis += rank + # In case data.shape is not statically known, we still can tell the slice is redundant if ends is sys.maxsize if ends_const.numpy().item() == _INT64_MAX: return True - if data.shape is None or data.shape.is_dynamic(axes_const.numpy().item()): + if data.shape is None or data.shape.is_dynamic(axis): logger.info("The value 'data' shape is not statically known.") return False - if ends_const.numpy().item() < data.shape[axes_const.numpy().item()]: + if ends_const.numpy().item() < data.shape[axis]: logger.info("The value 'end' is less than the shape of the specified axis.") return False @@ -78,7 +89,15 @@ def _potential_redundant_slice(op, data, starts, ends, axes, steps): return op.Slice(data, starts, ends, axes, steps, _outputs=["slice_output"]) -def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_): +def _same_shape( + op, + data: ir.Value, + slice_output: ir.Value, + steps: ir.Value, + starts: ir.Value | None = None, + axes: ir.Value | None = None, + **_, +) -> bool: """Check if the shape of the slice output is the same as the data.""" if data.shape is None or slice_output.shape is None: return False @@ -91,6 +110,20 @@ def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_ elif not _ir_utils.is_singleton_value(steps, 1): return False + if starts is not None: + starts_np = _ir_utils.get_numpy_value(starts) + if starts_np is not None and not all(s == 0 for s in starts_np.flat): + return False + + if axes is not None and data.shape is not None: + rank = data.shape.rank() + if rank is not None: + axes_np = _ir_utils.get_numpy_value(axes) + if axes_np is not None: + for a in axes_np.flat: + if a < -rank or a >= rank: + return False + return _ir_utils.same_shape(data.shape, slice_output.shape) From 70ccf7ffe19878df45d6b318f9b66138063e11d3 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Thu, 5 Nov 2026 09:47:43 +0530 Subject: [PATCH 3/6] fix(rewriter): normalize and clamp negative slice starts in same_shape onnx-ir 1.1.0 guarantees that Shape.rank() returns an int, and Shape indexing and is_dynamic() natively accept negative axis indices. Remove redundant rank is not None guards and explicit axis += rank normalization while retaining axis bounds checks. Additionally, update _same_shape to normalize and clamp negative slice starts against static dimension bounds instead of strictly checking for zero. If effective starts evaluate to 0 and all steps are 1 with matching output shapes, the slice can be collapsed as an identity operation. Reject non-zero starts on dynamic dimensions or when effective start != 0. Signed-off-by: Hardik Kaurani --- .../rewriter/rules/common/_collapse_slices.py | 48 +++++++++++++------ 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices.py b/onnxscript/rewriter/rules/common/_collapse_slices.py index 2a2d36b313..7f45997de3 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices.py @@ -59,12 +59,9 @@ def _check_if_redundant_slice( axis = int(axes_const.numpy().item()) if data.shape is not None: rank = data.shape.rank() - if rank is not None: - if axis < -rank or axis >= rank: - logger.info("The value 'axis' is out of range for the data rank.") - return False - if axis < 0: - axis += rank + if axis < -rank or axis >= rank: + logger.info("The value 'axis' is out of range for the data rank.") + return False # In case data.shape is not statically known, we still can tell the slice is redundant if ends is sys.maxsize if ends_const.numpy().item() == _INT64_MAX: @@ -110,18 +107,39 @@ def _same_shape( elif not _ir_utils.is_singleton_value(steps, 1): return False + rank = data.shape.rank() + + axes_np = _ir_utils.get_numpy_value(axes) if axes is not None else None + if axes_np is not None: + for a in axes_np.flat: + if a < -rank or a >= rank: + return False + if starts is not None: starts_np = _ir_utils.get_numpy_value(starts) - if starts_np is not None and not all(s == 0 for s in starts_np.flat): - return False - - if axes is not None and data.shape is not None: - rank = data.shape.rank() - if rank is not None: - axes_np = _ir_utils.get_numpy_value(axes) + if starts_np is not None: if axes_np is not None: - for a in axes_np.flat: - if a < -rank or a >= rank: + axes_list = [int(a) for a in axes_np.flat] + if len(axes_list) != len(starts_np.flat): + return False + elif axes is None: + axes_list = list(range(len(starts_np.flat))) + else: + if not all(s == 0 for s in starts_np.flat): + return False + axes_list = [] + + for s, a in zip(starts_np.flat, axes_list): + if a < -rank or a >= rank: + return False + if data.shape.is_dynamic(a): + if s != 0: + return False + else: + dim = data.shape[a] + eff_start = s + dim if s < 0 else s + eff_start = max(0, min(dim, eff_start)) + if eff_start != 0: return False return _ir_utils.same_shape(data.shape, slice_output.shape) From 648d4adc25be40e824ead5accb6c648d3abb4b10 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Thu, 5 Nov 2026 09:48:56 +0530 Subject: [PATCH 4/6] test(rewriter): add regression tests for negative and clamped starts collapsing to identity Add regression tests verifying that a Slice with negative starts that normalize (e.g. starts=[-8] on dimension 8) or clamp (e.g. starts=[-10] on dimension 8) to 0 along with matching shapes and unit steps are recognized as identity slices and collapsed into Identity nodes with verified numerical parity. Signed-off-by: Hardik Kaurani --- .../rules/common/_collapse_slices_test.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices_test.py b/onnxscript/rewriter/rules/common/_collapse_slices_test.py index 06d11654d7..9ad4de3d10 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices_test.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices_test.py @@ -256,3 +256,51 @@ def test_slice_valid_negative_axis_collapses(self): model, (np.random.rand(4, 8).astype(np.float32),), ) + + def test_slice_negative_normalized_start_collapses(self): + """Slice with negative start normalizing to 0 (e.g., -8 on dim 8) should collapse.""" + model_text = """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertIn("Identity", [node.op_type for node in model.graph]) + testing.assert_numerically_equal( + original, + model, + (np.random.rand(4, 8).astype(np.float32),), + ) + + def test_slice_negative_clamped_start_collapses(self): + """Slice with negative start clamping to 0 (e.g., -10 on dim 8) should collapse.""" + model_text = """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + original = ir.from_onnx_text(model_text) + model = ir.from_onnx_text(model_text) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertIn("Identity", [node.op_type for node in model.graph]) + testing.assert_numerically_equal( + original, + model, + (np.random.rand(4, 8).astype(np.float32),), + ) From cf0f614c723c137d174d85222df0728b7f24c396 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Thu, 5 Nov 2026 09:49:37 +0530 Subject: [PATCH 5/6] test(rewriter): add regression test for multi-axis nonzero start rejection in same_shape Add a regression test for multi-axis Slice with matching input/output shape annotations [4, 8], unit steps [1, 1], valid axes [0, 1], and nonzero starts [0, 1], asserting that no rewrite occurs in _same_shape. Signed-off-by: Hardik Kaurani --- .../rules/common/_collapse_slices_test.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices_test.py b/onnxscript/rewriter/rules/common/_collapse_slices_test.py index 9ad4de3d10..1deb9b3af1 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices_test.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices_test.py @@ -304,3 +304,21 @@ def test_slice_negative_clamped_start_collapses(self): model, (np.random.rand(4, 8).astype(np.float32),), ) + + def test_slice_multi_axis_nonzero_start_does_not_collapse(self): + """Slice with matching shapes, unit steps, valid axes, but nonzero start should not collapse.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) From f41d6b8ecdb6dc82850373b49a26bb067f2d59d9 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Thu, 5 Nov 2026 09:50:11 +0530 Subject: [PATCH 6/6] test(rewriter): add regression tests for single-axis nonzero start and multi-axis axis bounds in same_shape Add regression tests verifying that a single-axis Slice with nonzero start on matching shapes does not collapse, and that a multi-axis Slice containing an out-of-bounds axis does not collapse in _same_shape. Signed-off-by: Hardik Kaurani --- .../rules/common/_collapse_slices_test.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/onnxscript/rewriter/rules/common/_collapse_slices_test.py b/onnxscript/rewriter/rules/common/_collapse_slices_test.py index 1deb9b3af1..520bab1228 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices_test.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices_test.py @@ -322,3 +322,39 @@ def test_slice_multi_axis_nonzero_start_does_not_collapse(self): ) count = _collapse_slices.rules.apply_to_model(model) self.assertEqual(count, 0) + + def test_slice_single_axis_nonzero_start_does_not_collapse(self): + """Slice with matching shapes, unit steps, valid axis, but nonzero start should not collapse.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0) + + def test_slice_multi_axis_out_of_bounds_axis_does_not_collapse(self): + """Multi-axis slice containing an out-of-bounds axis should not collapse.""" + model = ir.from_onnx_text( + """ + + agraph (float[4, 8] data) => (float[4, 8] output) + { + starts = Constant() + ends = Constant() + axes = Constant() + steps = Constant() + output = Slice (data, starts, ends, axes, steps) + } + """ + ) + count = _collapse_slices.rules.apply_to_model(model) + self.assertEqual(count, 0)