Rewriter: validate and normalize axis in collapse_slices

onnxscript/rewriter/rules/common/_collapse_slices.py

@@ -55,13 +55,21 @@ 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 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:
         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 +86,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 +107,41 @@ 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
 
+    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:
+            if axes_np is not None:
+                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)
 
 

onnxscript/rewriter/rules/common/_collapse_slices_test.py

@@ -178,3 +178,183 @@ 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(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[1] {0}>()
+                ends = Constant<value: tensor = int64[1] {10}>()
+                axes = Constant<value: tensor = int64[1] {5}>()
+                steps = Constant<value: tensor = int64[1] {1}>()
+                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(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[1] {0}>()
+                ends = Constant<value: tensor = int64[1] {10}>()
+                axes = Constant<value: tensor = int64[1] {-5}>()
+                steps = Constant<value: tensor = int64[1] {1}>()
+                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"""
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {{
+                starts = Constant<value: tensor = int64[1] {{0}}>()
+                ends = Constant<value: tensor = int64[1] {{{_INT64_MAX}}}>()
+                axes = Constant<value: tensor = int64[1] {{5}}>()
+                steps = Constant<value: tensor = int64[1] {{1}}>()
+                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 = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[1] {0}>()
+                ends = Constant<value: tensor = int64[1] {10}>()
+                axes = Constant<value: tensor = int64[1] {-1}>()
+                steps = Constant<value: tensor = int64[1] {1}>()
+                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_normalized_start_collapses(self):
+        """Slice with negative start normalizing to 0 (e.g., -8 on dim 8) should collapse."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[1] {-8}>()
+                ends = Constant<value: tensor = int64[1] {8}>()
+                axes = Constant<value: tensor = int64[1] {1}>()
+                steps = Constant<value: tensor = int64[1] {1}>()
+                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 = """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[1] {-10}>()
+                ends = Constant<value: tensor = int64[1] {8}>()
+                axes = Constant<value: tensor = int64[1] {1}>()
+                steps = Constant<value: tensor = int64[1] {1}>()
+                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_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(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[2] {0, 1}>()
+                ends = Constant<value: tensor = int64[2] {4, 8}>()
+                axes = Constant<value: tensor = int64[2] {0, 1}>()
+                steps = Constant<value: tensor = int64[2] {1, 1}>()
+                output = Slice (data, starts, ends, axes, steps)
+            }
+        """
+        )
+        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(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[1] {1}>()
+                ends = Constant<value: tensor = int64[1] {8}>()
+                axes = Constant<value: tensor = int64[1] {1}>()
+                steps = Constant<value: tensor = int64[1] {1}>()
+                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(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[4, 8] data) => (float[4, 8] output)
+            {
+                starts = Constant<value: tensor = int64[2] {0, 0}>()
+                ends = Constant<value: tensor = int64[2] {4, 8}>()
+                axes = Constant<value: tensor = int64[2] {0, 5}>()
+                steps = Constant<value: tensor = int64[2] {1, 1}>()
+                output = Slice (data, starts, ends, axes, steps)
+            }
+        """
+        )
+        count = _collapse_slices.rules.apply_to_model(model)
+        self.assertEqual(count, 0)