Add remaining high-priority rewriter extended tests#2899
Merged
gramalingam merged 5 commits intomainfrom Apr 23, 2026
Merged
Conversation
…les) Extended test coverage for: - CastConstantOfShape: dtype variants (int32, float64, bfloat16), same-dtype positive - FuseConvAffine: non-constant weight/bias/scale negatives, padded pre-conv negative - RedundantScatterND: axis=0 dynamic positive, partial indices negative, shape mismatch negative - PartialRotaryEmbedding23: positive with dim attr validation, boundary mismatch negative, already-has-attr negative - GQA (rules/fusion): wrong group count negative, different head config positive - GQA (ort_fusions): no RotaryEmbedding negative - GQA packed QKV: misaligned slice boundaries negative - cos_sin_cache: non-constant inv_freq negative, constant inv_freq positive with ORT validation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Also adds a positive test with numerical validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…alue_info Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add _make_partial_rotary_script(mismatched) to _rotary_embedding_models.py that uses a traced-if to generate matching or mismatched slice boundaries. This replaces the graph-mutation approach in the test with a clean model built directly from @script(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR expands the ONNXScript rewriter “extended” test suite to cover additional high-priority fusion/rewrite gaps, and introduces a small model-generator helper to build cleaner negative rotary-embedding cases.
Changes:
- Added new extended test modules for multiple rewriter rules (fusion, common, and ort_fusions).
- Added negative/positive scenarios including boundary mismatches and prerequisite-missing cases.
- Extended the rotary-embedding model helpers with a script generator for partial rotary patterns.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxscript/rewriter/rules/fusion/_rotary_embedding_extended_test.py | Adds extended tests for PartialRotaryEmbedding23 (positive + negative cases). |
| onnxscript/rewriter/rules/fusion/_gqa_extended_test.py | Adds extended tests for fusion GQA rule (negative wrong-group + a positive alt config). |
| onnxscript/rewriter/rules/common/_redundant_scatter_nd_extended_test.py | Adds extended tests for redundant ScatterND elimination. |
| onnxscript/rewriter/rules/common/_fuse_conv_affine_extended_test.py | Adds extended tests for Conv/Affine fusion rules (positive + multiple negatives). |
| onnxscript/rewriter/rules/common/_cast_constant_of_shape_extended_test.py | Adds extended dtype-coverage tests for Cast+ConstantOfShape fusion. |
| onnxscript/rewriter/ort_fusions/gqa_packed_qkv_extended_test.py | Adds extended negative test for packed-QKV GQA fusion. |
| onnxscript/rewriter/ort_fusions/gqa_extended_test.py | Adds an extended negative test for ort_fusions GQA (no RotaryEmbedding prerequisite). |
| onnxscript/rewriter/ort_fusions/cos_sin_cache_extended_test.py | Adds extended tests for cos/sin cache fusion (dynamic inv_freq negative + sanity positive). |
| onnxscript/rewriter/models/_rotary_embedding_models.py | Adds _make_partial_rotary_script(mismatched=...) generator to support negative test construction. |
- Fix mixed 'import onnxscript' + 'from onnxscript import ...' in 5 files by using only 'from onnxscript import ...' style - Fix docstrings in gqa_extended_test.py and _redundant_scatter_nd_extended_test.py to match actual test coverage - Use symbolic dim N in scatter_nd dynamic test to prevent constant folding from eliminating the pattern nodes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
approved these changes
Apr 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add 22 extended tests across 8 test files covering the remaining high-priority gaps identified in the rewriter test gap analysis.
New test files
_cast_constant_of_shape_extended_test.py_cast_constant_of_shape_fuse_conv_affine_extended_test.py_fuse_conv_affine_redundant_scatter_nd_extended_test.py_redundant_scatter_nd_rotary_embedding_extended_test.py(rules/fusion)PartialRotaryEmbedding23Fusion_gqa_extended_test.py(rules/fusion)GroupQueryAttentionFusiongqa_extended_test.py(ort_fusions)gqaORT fusiongqa_packed_qkv_extended_test.py(ort_fusions)gqapacked QKVcos_sin_cache_extended_test.py(ort_fusions)cos_sin_cacheKey patterns used
@script()API withop.Constant(value=...)for embedded constants"B","S") where appropriateModified existing file
_rotary_embedding_models.py: Added_make_partial_rotary_script(mismatched)generator for clean negative test constructionFollow-up to PR #2896.