Add remaining high-priority rewriter extended tests by gramalingam · Pull Request #2899 · microsoft/onnxscript

gramalingam · 2026-04-22T05:40:09Z

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

File	Tests	Rule module
`_cast_constant_of_shape_extended_test.py`	5	`_cast_constant_of_shape`
`_fuse_conv_affine_extended_test.py`	5	`_fuse_conv_affine`
`_redundant_scatter_nd_extended_test.py`	3	`_redundant_scatter_nd`
`_rotary_embedding_extended_test.py` (rules/fusion)	3	`PartialRotaryEmbedding23Fusion`
`_gqa_extended_test.py` (rules/fusion)	2	`GroupQueryAttentionFusion`
`gqa_extended_test.py` (ort_fusions)	1	`gqa` ORT fusion
`gqa_packed_qkv_extended_test.py` (ort_fusions)	1	`gqa` packed QKV
`cos_sin_cache_extended_test.py` (ort_fusions)	2	`cos_sin_cache`

Key patterns used

@script() API with op.Constant(value=...) for embedded constants
Symbolic dimensions ("B", "S") where appropriate
Numerical validation via ORT or ONNX reference implementation
Negative tests: non-constant operands, padded convolutions, missing prerequisites, mismatched boundaries
Script generator with traced-if for clean negative test model construction

Modified existing file

_rotary_embedding_models.py: Added _make_partial_rotary_script(mismatched) generator for clean negative test construction

Follow-up to PR #2896.

…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>

+import onnx_ir as ir
+from packaging import version
+
+import onnxscript


+import numpy
+import onnx_ir as ir
+
+import onnxscript


+import onnx_ir as ir
+import onnx_ir.passes.common.shape_inference as shape_inference
+
+import onnxscript


+import onnx_ir as ir
+import onnx_ir.passes.common.shape_inference as shape_inference
+
+import onnxscript


codecov · 2026-04-22T05:43:54Z

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
13761	1	13760	2641

View the top 1 failed test(s) by shortest run time

onnxscript.rewriter.ort_fusions.mha_test.TestMultiHeadAttention::test_smollm

Stack Traces | 3.29s run time

onnxscript\rewriter\ort_fusions\mha_test.py:23: in test_smollm
    model = smollm_test.get_onnx_model()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\rewriter\models\_smollm_2.py:453: in get_onnx_model
    model_proto = make_model_with_random_weights()
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\rewriter\models\_smollm_2.py:416: in make_model_with_random_weights
    lm_head_weight = numpy.random.rand(49152, 2048).astype(numpy.float32)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   numpy._core._exceptions._ArrayMemoryError: Unable to allocate 384. MiB for an array with shape (49152, 2048) and data type float32

View the full list of 1 ❄️ flaky test(s)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logsumexp_cpu_float16

Flake rate in main: 14.50% (Passed 2948 times, Failed 500 times)

Stack Traces | 0.636s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 1 / 5 (20.0%)
E   Greatest absolute difference: 2.288818359375e-05 at index (1,) (up to 1e-05 allowed)
E   Greatest relative difference: 0.0022869110107421875 at index (1,) (up to 0.001 allowed)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

Copilot

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.

Copilot · 2026-04-22T06:00:03Z

+        def model_fn(data: FLOAT[8, 16], updates: FLOAT[8, 16]) -> FLOAT[8, 16]:
+            axis = op.Constant(value_int=0)
+            shape = op.Shape(data, start=0)
+            dim = op.Gather(shape, axis, axis=0)
+            full_range = op.Range(0, dim, 1)
+            full_range_2d = op.Unsqueeze(full_range, [-1])
+            scattered = op.ScatterND(data, full_range_2d, updates, reduction="none")
+            return scattered
+
+        model_proto = model_fn.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        onnx_check(model)
+        shape_inference(model)
+        onnxscript.optimizer.fold_constants(model)
+        count = _redundant_scatter_nd.rules.apply_to_model(model)
+        self.assertEqual(count, 1)


In test_dynamic_indices_axis_0, the input types are fully static (FLOAT[8, 16]) and the test runs fold_constants() before applying the rewrite. With fully-known shapes, constant folding can precompute Shape/Gather/Range/Unsqueeze, removing the dynamic-indices subgraph that ScatterAllDynamic matches (and ScatterAllStatic may not match due to reduction="none"). Consider making the updated axis dimension symbolic (e.g., FLOAT["N", 16]) or skipping constant-folding here so the dynamic-pattern nodes remain and the rule is actually exercised reliably.

Copilot · 2026-04-22T06:00:04Z

+Adds coverage for: different axis (dynamic), partial scatter (negative),
+shape mismatch (negative for static), and reduction="add" (negative for static).


The module docstring claims coverage for reduction="add" and an additional negative case, but this file currently only includes the axis=0 positive test plus two static negative tests (partial indices, shape mismatch). Either add the missing test(s) or update the docstring to reflect the actual coverage.

Suggested change

Adds coverage for: different axis (dynamic), partial scatter (negative),

shape mismatch (negative for static), and reduction="add" (negative for static).

Adds coverage for: dynamic full-range scatter on axis=0 (positive),

partial scatter (negative for static), and shape mismatch (negative for static).

Copilot · 2026-04-22T06:00:04Z

+
+"""Extended tests for GQA ort_fusions fusion.
+
+Adds coverage for: wrong expand shape (negative — group count mismatch).


The file docstring says this adds coverage for a “wrong expand shape (negative — group count mismatch)”, but the only test in this file is test_no_rotary_embedding_no_gqa_fusion (no RotaryEmbedding prerequisite). Please update the docstring (or rename the test/add the intended group-mismatch case) so the stated coverage matches what the test actually validates.

Suggested change

Adds coverage for: wrong expand shape (negative — group count mismatch).

Adds coverage for: missing RotaryEmbedding (negative — no GQA fusion).

gramalingam and others added 4 commits April 22, 2026 03:56

Rewrite conv-affine tests to use @script() instead of ir.tape API

d2b66c6

Also adds a positive test with numerical validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Use symbolic dim 'S' instead of concrete value in GQA extended test v…

d16b063

…alue_info Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

github-project-automation Bot added this to ONNX Script Review Board Apr 22, 2026

github-project-automation Bot moved this to Todo in ONNX Script Review Board Apr 22, 2026

github-code-quality Bot found potential problems Apr 22, 2026

View reviewed changes

justinchuby requested a review from Copilot April 22, 2026 05:55

Copilot started reviewing on behalf of justinchuby April 22, 2026 05:56 View session

Copilot AI reviewed Apr 22, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add remaining high-priority rewriter extended tests#2899

Add remaining high-priority rewriter extended tests#2899
gramalingam wants to merge 4 commits intomainfrom
rama/tests2

gramalingam commented Apr 22, 2026

Uh oh!

codecov Bot commented Apr 22, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 22, 2026

Uh oh!

Copilot AI Apr 22, 2026

Uh oh!

Copilot AI Apr 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		Adds coverage for: different axis (dynamic), partial scatter (negative),
		shape mismatch (negative for static), and reduction="add" (negative for static).


		"""Extended tests for GQA ort_fusions fusion.

		Adds coverage for: wrong expand shape (negative — group count mismatch).

	Adds coverage for: wrong expand shape (negative — group count mismatch).
	Adds coverage for: missing RotaryEmbedding (negative — no GQA fusion).

Conversation

gramalingam commented Apr 22, 2026

Summary

New test files

Key patterns used

Modified existing file

Uh oh!

codecov Bot commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ 1 Tests Failed:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Apr 22, 2026 •

edited

Loading