[Hardik K (hardikkaurani)]

Summary

This PR adds comprehensive validation to the _cast_constant_of_shape rewrite rules to prevent invalid or unsafe Cast fusion into ConstantOfShape.

Specifically:

• Rejects unsupported STRING and UNDEFINED target data types. Under the ONNX specification, ConstantOfShape does not support STRING output element types; preventing this fusion preserves model validity and prevents downstream serialization errors.
• Rejects reference-valued attributes for both target dtype (dtype.is_ref()) and scalar fill values (scalar.is_ref()), ensuring unexpanded function attributes are preserved.
• Validates that explicit fill value attributes are single-element scalar tensors (scalar_tensor.size == 1) and not string tensors.
• Restricts the no-value target pattern with _allow_other_attributes=False so ConstantOfShape nodes carrying an explicit value attribute cannot match the fallback rule and be silently replaced with zero.
• Removes unrelated PyTorch nightly OpInfo test workarounds per maintainer guidance, keeping this PR strictly focused on the ConstantOfShape rewrite fix.

Problem

Previously, cast_constant_of_shape_rule and cast_constant_of_shape_without_value_rule lacked sufficient validation:

1. STRING and UNDEFINED Target Types: The ONNX specification does not support STRING or UNDEFINED output tensor element types for ConstantOfShape. Fusing a Cast targeting ir.DataType.STRING results in an invalid ONNX model and causes serialization failures (SerdeError: Bitwidth not available for ONNX data type: STRING).
2. Reference-Valued Attributes: Reference attributes (such as @target or @fill within ONNX functions) do not possess concrete values at rewrite time and cannot be evaluated statically.
3. Non-Scalar Fill Values: The fusion logic requires a single-element scalar fill tensor. Non-singleton tensors cannot be fused into scalar fill values.
4. Unsafe Fallback Matching: The no-value pattern previously permitted other attributes by default. Consequently, a ConstantOfShape node with an explicit value attribute that failed the primary rule's guards could fall through to the fallback rule, silently discarding its original fill value in favor of zero.

Solution

• Updated check_cast_constant_of_shape and check_cast_constant_of_shape_without_value to reject ir.DataType.STRING and ir.DataType.UNDEFINED.
• Added guards against reference attributes on both the target dtype and the fill value tensor.
• Added validation ensuring fill tensors have size == 1 and non-string dtype, using scalar.as_tensor() for explicit tensor access.
• Constrained the fallback pattern using op.ConstantOfShape(shape, _allow_other_attributes=False).
• Reverted unrelated nightly CI test xfail workarounds to keep this PR focused entirely on the rewriter fix.

Tests

Added and updated regression tests in onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py:

• test_cast_after_constant_of_shape_is_fused: Validates legitimate fusion with an explicit fill value.
• test_cast_after_constant_of_shape_without_value_is_fused: Validates legitimate fusion without an explicit fill value.
• test_cast_to_string_after_constant_of_shape_is_not_fused: Verifies that STRING-target fusion is rejected with an explicit fill value because ConstantOfShape does not support STRING outputs.
• test_cast_to_string_after_constant_of_shape_without_value_is_not_fused: Verifies that STRING-target fusion is rejected without an explicit fill value because ConstantOfShape does not support STRING outputs.
• test_cast_after_constant_of_shape_with_ref_attr_is_not_fused: Verifies that reference-valued fill attributes (@fill) are preserved.
• test_cast_with_ref_attr_to_after_constant_of_shape_is_not_fused: Verifies that reference-valued target dtype attributes (@target) are preserved with an explicit fill value.
• test_cast_with_ref_attr_to_after_constant_of_shape_without_value_is_not_fused: Verifies that reference-valued target dtype attributes (@target) are preserved without an explicit fill value.
• test_cast_to_undefined_after_constant_of_shape_is_not_fused: Verifies that UNDEFINED-target fusion is rejected with an explicit fill value.
• test_cast_to_undefined_after_constant_of_shape_without_value_is_not_fused: Verifies that UNDEFINED-target fusion is rejected without an explicit fill value.
• test_cast_after_constant_of_shape_with_non_scalar_value_is_not_fused: Verifies that non-scalar fill tensors (float[2]) are rejected.
• test_cast_after_constant_of_shape_with_string_value_is_not_fused: Verifies that string-valued ConstantOfShape fill tensors are rejected.
• test_cast_after_constant_of_shape_without_value_numerical_correctness: Verifies numerical equivalence for legitimate no-value fusion.

[Hardik K (hardikkaurani)]

Hi Ti-Tai Wang (@titaiwangms), Could you please review this PR when you have a moment? It adds condition checks to _cast_constant_of_shape to prevent invalid string tensor generation and subsequent serialization failures. 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 fallback pattern can match ConstantOfShape nodes with an explicit value and silently replace that fill value with zero, requiring stricter attribute matching and a regression test.

1 open issue

• Prevent no-value fallback from overriding explicit fill values · New

_{Review effort: Balanced}

[Hardik K (hardikkaurani)]

Addressed in the latest commits. Added _allow_other_attributes=False to op.ConstantOfShape(shape, _allow_other_attributes=False) in the no-value target pattern so that nodes with explicit value attributes do not match the fallback. Also added test_cast_after_constant_of_shape_with_ref_attr_is_not_fused using the full rule set to ensure reference attribute @fill is preserved without being overwritten with zero.

[Ti-Tai Wang (titaiwangms)]

I recommend changes before merging: the new reference-value guard can fall through to the no-value rule and produce an incorrect zero-filled result.

In this rewriter, patterns allow extra attributes by default. Consequently, op.ConstantOfShape(shape) also matches a ConstantOfShape with an explicit value attribute. For a function containing ConstantOfShape(value=@fill) -> Cast(to=FLOAT16), the first rule now correctly rejects the reference attribute, but the second rule still matches and checks only the concrete target dtype. Its replacement discards @fill and constructs a zero-filled ConstantOfShape. A nonzero fill value therefore becomes a silent numerical error instead of being left unchanged.

Minimal fix: use op.ConstantOfShape(shape, _allow_other_attributes=False) in the no-value pattern, and add a regression test with a reference value and a concrete numeric Cast target. The two new STRING-target tests do not exercise this fallback path.

The STRING guard itself is correct: ConstantOfShape does not support STRING outputs, regardless of whether serialization supports string tensors. The use of IR attribute/type/size checks is reasonable; scalar.as_tensor() would also make the tensor access more explicit than .value.

Validation: static review of the diff, pattern-matcher defaults, rule fallthrough/function traversal, and ONNX ConstantOfShape specification only; tests were not run locally because onnx_ir is unavailable in this environment.

[Hardik K (hardikkaurani)]

Hi Ti-Tai Wang (@titaiwangms),

Thank you for the review and catching that fallback issue! I have updated the PR accordingly across three focused commits:

1. Restrict no-value pattern: Added _allow_other_attributes=False to op.ConstantOfShape(shape, _allow_other_attributes=False) in cast_constant_of_shape_without_value. This prevents ConstantOfShape nodes carrying an explicit value attribute (such as reference attributes or non-scalar tensors rejected by the first rule) from falling through to the no-value rule and being replaced with zero. Also switched to scalar.as_tensor() for explicit tensor access.
2. Reference attribute regression test: Added test_cast_after_constant_of_shape_with_ref_attr_is_not_fused using the complete rule set, verifying that a function containing ConstantOfShape <value: tensor = @fill> (shape) -> Cast <to = 1> is not rewritten and preserves the reference attribute intact.
3. Non-scalar & numerical tests: Added test_cast_after_constant_of_shape_with_non_scalar_value_is_not_fused to verify rejection of non-singleton tensors, along with test_cast_after_constant_of_shape_without_value_numerical_correctness verifying numerical equivalence for valid no-value fusion.

All 7 tests in _cast_constant_of_shape_test.py as well as ruff formatting and lint checks pass cleanly.

[Ti-Tai Wang (titaiwangms)]

The explicit-value fallback issue is resolved by _allow_other_attributes=False, and the added reference-value regression test covers the previously unsafe path. The use of as_tensor() is also clearer. No blocking findings remain in my static review.

Please update the PR description to include the strict no-value pattern matching and the new reference-value/non-singleton regression tests. Please also clarify that STRING fusion is invalid because ConstantOfShape does not support STRING output, not only because of serialization behavior.

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

[codecov]

Codecov Report

❌ Patch coverage is 94.05941% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.41%. Comparing base (682b59d) to head (ea06132).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
...t/rewriter/rules/common/_cast_constant_of_shape.py	78.57%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3145      +/-   ##
==========================================
+ Coverage   73.29%   73.41%   +0.12%     
==========================================
  Files         267      269       +2     
  Lines       33214    33662     +448     
  Branches     3183     3218      +35     
==========================================
+ Hits        24344    24713     +369     
- Misses       7801     7876      +75     
- Partials     1069     1073       +4     

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

[Ti-Tai Wang (titaiwangms)]

The original ConstantOfShape rewrite fix still looks good, but I do not recommend merging the latest revision (ea06132) with the newly added nightly workarounds as-is. My earlier approval covered commit 6c5329a, before these additional changes.

Blocking CI issue: In tests/function_libs/torch_lib/ops_test_data.py, lines 510–519 and 614–623, the new addbmm/baddbmm xfail matchers do not restrict dtype. They also match working int32 and int64 samples, causing the expected-failure mechanism to fail those tests with “Test unexpectedly passed.” The latest Linux torch-nightly job shows exactly four such failures: addbmm and baddbmm, each for int32 and int64. The macOS and Windows torch-nightly jobs are also failing.

Evidence: https://github.com/microsoft/onnxscript/actions/runs/42106844617/job/127538046524 (the run head SHA matches ea06132).

Please preferably move these unrelated nightly workarounds into a separate PR so this PR remains focused on the ConstantOfShape fix. If they remain here, narrow the xfails to the exact dtypes and samples demonstrated to fail, and rerun the affected tests. Please do not replace them with broad skips.

The description now covers the intended rewrite changes, but it contains control characters that corrupt occurrences of value, operator names, and test names. Please clean those up as well.

Validation: reviewed the added commit and the latest CI failure logs; no local tests were run.

[Hardik K (hardikkaurani)]

Hi Ti-Tai Wang (@titaiwangms),

Thank you for the guidance. I've updated the branch with commit (754c4f86888e0e2c492fdda317b2a46605901e32)](https://github.com/microsoft/onnxscript/commit/754c4f86888e0e2c492fdda317b2a46605901e32) to address all the feedback:

1. Removed Unrelated Nightly Workarounds: Completely reverted tests/function_libs/torch_lib/ops_test_data.py to match upstream/main, removing the nightly addbmm, baddbmm, and logit xfail matchers. The PR diff now exclusively touches _cast_constant_of_shape.py and its dedicated test suite, _cast_constant_of_shape_test.py.

2. Clarified STRING Output Rationale: Updated the docstrings and PR description to explicitly state that ConstantOfShape does not support STRING output element types under the ONNX specification, rather than focusing solely on serialization behavior.

3. Cleaned PR Description: Refreshed the PR description with clean, standard formatting, removing corrupted characters while accurately documenting the strict _allow_other_attributes=False matching, reference-attribute checks, non-singleton scalar validation, and all 12 regression test cases.

All code and test formatting checks pass cleanly via Ruff.

Please let me know if any further adjustments are needed. Thank you!