[Hardik K (hardikkaurani)]

Summary

This PR hardens the _collapse_slices rewrite rule against invalid axis values that could cause unhandled IndexError crashes or incorrect identity rewrites.

It adds axis bounds validation and improves same-shape detection to account for ONNX Slice start normalization and clamping, preserving valid identity-slice optimizations.

Problem

In _collapse_slices._check_if_redundant_slice(), the axis value was extracted from the Slice node without validating whether it was within the valid range for the input tensor.

This caused two issues:

1. Out-of-bounds axes: Invalid axes could cause data.shape.is_dynamic() to raise an unhandled IndexError. Additionally, when ends == _INT64_MAX, the rule could return True before validating the axis, incorrectly collapsing an invalid Slice into an Identity.

2. Incorrect same-shape detection: _same_shape() could treat matching input and output shapes as sufficient evidence for an identity rewrite without properly validating the slice parameters. Conversely, requiring every raw starts value to equal zero would reject valid identity slices whose negative starts normalize or clamp to zero.

Solution

Updated _collapse_slices with the following changes:

• Validate statically known axes against the input tensor's rank before performing the ends == _INT64_MAX early return.
• Reject out-of-bounds axes while preserving native negative-index support provided by IR Shape indexing.
• Improve _same_shape() to validate slice parameters rather than relying solely on matching input and output shapes.
• Calculate effective start positions for static dimensions using ONNX Slice normalization and clamping semantics.
• Allow valid negative starts that normalize or clamp to zero to retain identity-slice optimization.
• Reject nonzero effective starts and nonzero starts on dynamic dimensions when identity cannot be safely established.

These changes prevent invalid slices from being incorrectly rewritten while preserving valid identity optimizations.

Tests

Added and updated regression tests in _collapse_slices_test.py:

• test_slice_out_of_bounds_axis_does_not_crash_or_collapse: Verifies that a positive out-of-bounds axis is safely rejected.
• test_slice_negative_out_of_bounds_axis_does_not_crash_or_collapse: Verifies that a negative out-of-bounds axis is safely rejected.
• test_slice_out_of_bounds_axis_with_int64_max_ends_does_not_collapse: Ensures invalid axes are rejected even when ends == _INT64_MAX.
• test_slice_valid_negative_axis_collapses: Verifies valid negative-axis handling and numerical equivalence.
• test_slice_multi_axis_nonzero_start_does_not_collapse: Ensures a nonzero effective start does not produce an incorrect identity rewrite.
• Additional regression coverage verifies negative starts that normalize or clamp to zero and preserves numerical equivalence for valid identity slices.

[Hardik K (hardikkaurani)]

Hi Ti-Tai Wang (@titaiwangms), Could you please review this PR when you have a moment? It fixes an IndexError crash and prevents invalid axis values from being incorrectly collapsed into an Identity in _collapse_slices. All tests and lint checks are passing. Thank you for your time!

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

🟡 Changes recommended

The nonzero-start rejection lacks a regression test with matching shapes, unit steps, valid axes, and starts [0, 1] to ensure the _same_shape fallback does not incorrectly rewrite the Slice as an identity.

1 open issue

• Add regression test for nonzero Slice start rejection · New

_{Review effort: Balanced}

[Ti-Tai Wang (titaiwangms)]

The axis-bounds checks fix a real issue: they prevent both an IndexError and the same-shape fallback from eliminating a Slice with a statically known invalid axis.

One adjustment I recommend before merging: the new requirement in _same_shape that every raw starts value equal zero is unnecessarily restrictive. For example, with data.shape=[4,8], starts=[-8], ends=[8], axes=[1], and steps=[1], ONNX Slice normalizes the start to zero, so this is a valid identity slice. The previous same-shape rule could eliminate it; the new check prevents that optimization. Please either keep this PR focused on axis validation or account for Slice start normalization/clamping and add a regression test for this case.

Regarding IR reuse: in the minimum supported onnx-ir version (1.1.0), Shape.rank() always returns an integer, and both shape indexing and is_dynamic() already support valid negative indices. The rank is not None branches and explicit negative-axis normalization are therefore unnecessary here. The substantive fix is bounds validation, not new negative-axis support.

Validation: static review of the diff, matcher behavior, IR implementation, and ONNX Slice specification only; tests were not run locally because onnx_ir is unavailable in this environment.

[Hardik K (hardikkaurani)]

Thanks for the review and suggestions! Ti-Tai Wang (@titaiwangms)

I have updated the PR accordingly across four focused commits:

1. Slice start normalization and clamping: Updated _same_shape to compute effective starts against static dimension bounds (eff_start = max(0, min(dim, s + dim if s < 0 else s))). Negative starts that normalize or clamp to 0 (such as starts=[-8] or starts=[-10] on dimension 8) now correctly collapse to identity when output shapes and steps match. Nonzero starts on dynamic dimensions and non-zero effective starts are safely rejected.
2. IR simplification: Removed the redundant rank is not None checks and manual negative-axis normalization (axis += rank), relying on native negative index support in onnx-ir >= 1.1.0 while keeping axis bounds validation intact.
3. Regression tests: Added test cases covering negative normalized and clamped starts with numerical equivalence checks, multi-axis nonzero start rejection (starts=[0, 1]), single-axis nonzero start rejection, and multi-axis out-of-bounds axis handling.

All 17 tests in _collapse_slices_test.py as well as ruff check and ruff format pass cleanly.

[codecov]

Codecov Report

❌ Patch coverage is 81.39535% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.27%. Comparing base (9112df0) to head (f41d6b8).

Files with missing lines	Patch %	Lines
...nxscript/rewriter/rules/common/_collapse_slices.py	57.89%	8 Missing and 8 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3144      +/-   ##
==========================================
+ Coverage   73.25%   73.27%   +0.01%     
==========================================
  Files         269      269              
  Lines       33441    33524      +83     
  Branches     3192     3208      +16     
==========================================
+ Hits        24498    24565      +67     
- Misses       7874     7882       +8     
- Partials     1069     1077       +8     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[Ti-Tai Wang (titaiwangms)]

The revised changes address my previous concerns: negative starts are normalized/clamped correctly, and the regression tests cover the optimization and axis-validation cases. No blocking findings remain in my static review.

Please update the PR description to match the final implementation: negative-axis handling now relies on IR Shape indexing rather than explicit normalization, and the same-shape check accepts starts whose effective value normalizes/clamps to zero rather than requiring every raw start to be zero.

Approval is based on static review; local tests were not run. Please ensure the required CI checks pass before merging.

[Hardik K (hardikkaurani)]

Sure Ti-Tai Wang (@titaiwangms) !
Working on it.