You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Fix contracting dim match and batch broadcasting in broadcast_to_matmul - #3103
In _broadcast_to_matmul.py, check_if_not_need_reshape inspected reversed input dimensions with:
ifdim_from_anotin {1, dim_from_b}:
returnFalse
This introduced two issues:
Invalid graph rewrite on contracting dimension mismatch (idx == 0):
At idx == 0, dim_from_a and dim_from_b are the contracting dimensions $K$ of $A$ and $B$. In MatMul, $K$ cannot broadcast and must match exactly. Accepting dim_from_a == 1 when dim_from_b != 1 erroneously permitted rewriting valid reshape graphs into invalid MatMul nodes with mismatched contracting dimensions that fail ONNX shape inference.
Rejected rewrite for valid batch broadcasting (idx > 0):
At idx > 0, dim_from_a and dim_from_b are batch dimensions. Two dimensions are broadcastable if $d_A = d_B$, $d_A = 1$, or $d_B = 1$. The previous check rejected valid broadcasting when $d_B = 1$ and $d_A > 1$ (e.g. $[2, 3, 4]$ and $[1, 4, 5]$).
Changes
Enforce exact equality on contracting dimensions at idx == 0 (dim_from_a == dim_from_b).
Enable standard symmetric broadcasting on batch dimensions at idx > 0 (dim_from_a == dim_from_b or dim_from_a == 1 or dim_from_b == 1).
Validate removed reshapes preserve matrix computation and batch layout:
Validates that trailing 2-D matrix dimensions $(M, K)$ and $(K, N)$ match between the pre-reshape inputs and the reshaped tensors (accounting for 1-D vector expansion).
Validates that batch dimensions are not transposed or permuted across axes (batch shapes must match when right-aligned, allowing only insertion or removal of leading size-1 dimensions).
Unit tests added / updated in _broadcast_to_matmul_test.py:
test_reshape_matmul_reshape_replace_when_second_input_batch_dim_is_one: confirms redundant reshapes are removed when input $B$ has batch dimension 1.
test_reshape_matmul_reshape_remain_when_contracting_dimension_mismatches: element-count-preserving fixture confirming the rewrite is safely declined when contracting dimensions mismatch.
test_reshape_matmul_reshape_remain_when_reshapes_permute_batch_dimensions: confirms the rewrite is safely declined when reshapes alter batch-axis layout.
test_reshape_matmul_reshape_numerical_equivalence_when_second_input_batch_dim_is_one: end-to-end regression verifying numerical equivalence between original and optimized models with ONNX Runtime.
Validation
All 20 unit tests in _broadcast_to_matmul_test.py pass.
All 325 common rewriter rules tests pass (1 existing skip).
ruff check and ruff format --check pass with zero warnings or errors.
In _broadcast_to_matmul.py, check_if_not_need_reshape checked reversed dimensions using if dim_from_a not in {1, dim_from_b}:.
- At idx == 0 (the contracting dimension K), dim_from_a and dim_from_b must match exactly (dim_from_a == dim_from_b). Allowing dim_from_a == 1 when dim_from_b != 1 permitted rewriting graphs into invalid MatMul operations with incompatible contracting dimensions.
- At idx > 0 (the batch dimensions), two dimensions are broadcastable if dim_from_a == dim_from_b, dim_from_a == 1, or dim_from_b == 1. The previous check rejected valid broadcasting when dim_from_b == 1 and dim_from_a > 1.
This change enforces exact equality on contracting dimensions at idx == 0 and symmetric broadcasting on batch dimensions at idx > 0. Adds regression unit tests covering both cases.
I have submitted this pull request to resolve two correctness issues in _broadcast_to_matmul.py:
When contracting dimensions mismatch at idx == 0, exact equality is now required rather than allowing dim_from_a == 1, preventing invalid rewrites into incompatible MatMul operations.
When batch dimensions broadcast symmetrically at idx > 0, the check now correctly accounts for dim_from_b == 1 as well as dim_from_a == 1.
All 18 unit tests in _broadcast_to_matmul_test.py and all 322 common rewrite tests pass, and ruff checks and formatting are clean.
Please review when convenient and let me know if any adjustments are desired.
Following up on this PR when you have an opportunity. It fixes the contracting dimension check and adds symmetric batch dimension broadcasting in _broadcast_to_matmul.py. All 18 unit tests and rewrite suites pass cleanly. Please let me know if any adjustments are needed. Thank you!
Added validation in check_if_not_need_reshape to inspect the removed reshapes (shape_a and shape_b).
Verifies that the trailing 2-D matrix computation dimensions $(M, K)$ and $(K, N)$ match between pre-reshape inputs and reshaped tensors (accounting for 1-D vector expansion).
Verifies that batch dimensions are not permuted or transposed (they must match when right-aligned, permitting only leading size-1 dimensions to differ).
Added test_reshape_matmul_reshape_remain_when_reshapes_permute_batch_dimensions ensuring graphs with permuted batch dimensions are safely declined.
Added test_reshape_matmul_reshape_numerical_equivalence_when_second_input_batch_dim_is_one using ONNX Runtime to verify numerical equivalence.
Element-count-preserving fixture:
Updated test_reshape_matmul_reshape_remain_when_contracting_dimension_mismatches to use element-count-preserving shapes ([4, 3, 1] reshaped to [1, 3, 4], multiplied with [4, 4, 5]), ensuring the original graph is valid and verifying that the rewrite is safely declined due to contracting dimension mismatch.
All 20 unit tests in _broadcast_to_matmul_test.py and all 325 common rewriter rules tests pass cleanly. ruff check and ruff format are clean.
..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
return func(*(a + p.args), **p.kwargs, **kw)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
return exporter.export(model_onnx, function_name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
add(self._import_onnx_types(proto))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
ts = _translate_type(t.type)
^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E NotImplementedError: Unable to translate type optional_type {
E elem_type {
E tensor_type {
E elem_type: 1
E shape {
E dim {
E dim_value: 4
E }
E }
E }
E }
E }
E into onnxscript type.
..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
return func(*(a + p.args), **p.kwargs, **kw)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
return exporter.export(model_onnx, function_name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
add(self._import_onnx_types(proto))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
ts = _translate_type(t.type)
^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E NotImplementedError: Unable to translate type optional_type {
E elem_type {
E tensor_type {
E elem_type: 1
E shape {
E dim {
E dim_value: 4
E }
E }
E }
E }
E }
E into onnxscript type.
..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
return func(*(a + p.args), **p.kwargs, **kw)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
return exporter.export(model_onnx, function_name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
add(self._import_onnx_types(proto))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
ts = _translate_type(t.type)
^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E NotImplementedError: Unable to translate type optional_type {
E elem_type {
E tensor_type {
E elem_type: 1
E shape {
E dim {
E dim_value: 4
E }
E }
E }
E }
E }
E into onnxscript type.
..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
return func(*(a + p.args), **p.kwargs, **kw)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
return exporter.export(model_onnx, function_name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
add(self._import_onnx_types(proto))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
ts = _translate_type(t.type)
^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E NotImplementedError: Unable to translate type optional_type {
E elem_type {
E tensor_type {
E elem_type: 6
E shape {
E }
E }
E }
E }
E into onnxscript type.
The new symmetric batch predicate accepts (0, 1), but max(0, 1) predicts 1 instead of the correct broadcast dimension 0. This can transform a valid model into a graph with a different observable output shape, rather than merely skip an optimization.
A valid opset-17 counterexample, using default allowzero=0:
A: float[0,0,4]
B: float[1,4,5]
AR = Reshape(A, [0,0,4])
P = MatMul(AR, B) # [0,0,5]
Y = Reshape(P, [1,0,5]) # [1,0,5]; valid because both element counts are zero
The layout guard accepts the identity input Reshape. The current check predicts [1,0,5], matches the final target, and removes the Reshapes. The replacement MatMul(A,B) actually produces [0,0,5]. A downstream Shape(Y) consequently changes from [1,0,5] to [0,0,5].
Confirmed with both ONNX Runtime CPU (graph optimizations disabled) and ONNX ReferenceEvaluator, for both the one-input-Reshape and two-input-Reshape patterns. Each head rewrite count was 1; the reverse-reconstructed base rule declined both cases (count 0), using head infrastructure rather than a full base checkout.
Minimal fix: after checking compatibility, choose the non-singleton dimension instead of the numerical maximum, e.g. dim_from_b if dim_from_a == 1 else dim_from_a. Add an empty-tensor regression asserting the actual output shape or downstream Shape result, not only elementwise equality.
Non-blocking questions
The new layout guard also rejects matching, order-preserving batch flattening: A[2,3,4,5] -> [6,4,5], B[2,3,5,6] -> [6,5,6], followed by output reshape [2,3,4,6]. The original computation matches direct NumPy MatMul for the tested deterministic inputs, while the head declines a rewrite the reconstructed base accepts. Is this conservative reduction in optimization coverage intentional?
The target-shape resolver assumes allowzero=0, while the patterns do not constrain the attribute. Consider documenting/constraining this assumption or explicitly modeling allowzero=1. No separate valid-model miscompilation was established for this concern, so it is not a blocker.
- Select non-singleton batch dimension instead of max(dim_a, dim_b) when calculating MatMul broadcast shape to prevent incorrect predictions with 0-sized dimensions.
- Document allowzero=0 assumption in _resolve_target_shape.
- Add regression tests for empty-tensor zero-sized batch broadcasting for one-reshape and two-reshape patterns, verifying both rewrite acceptance/rejection and runtime output shape preservation.
Zero-sized batch broadcasting: Replaced max(dim_from_a, dim_from_b) with dim_from_b if dim_from_a == 1 else dim_from_a in _broadcast_to_matmul.py. When broadcasting a size-0 batch dimension with size 1, the broadcast dimension is now correctly predicted as 0 rather than 1, preventing the removal of output reshapes with different observable shapes.
Regression tests: Added empty-tensor tests in both TwoReshapesMatMulReshapeTest and OneReshapeMatMulReshapeTest confirming that:
When the target shape is [1, 0, 5], the rewrite is declined (count == 0) and ONNX Runtime confirms the observable runtime output shape is (1, 0, 5).
When the target shape matches [0, 0, 5], the rewrite is accepted (count == 1) and ONNX Runtime confirms the output shape is (0, 0, 5).
allowzero assumption: Documented the allowzero=0 assumption in the _resolve_target_shape docstring.
Non-blocking question 1 (batch flattening): The layout guard was kept intentionally conservative to strictly preserve batch-axis alignment and prevent any risk of transposing or altering batch dimensions.
Resolved shape_c against the matched MatMul output: Instead of treating shape_c as literal dimension values in check_if_not_need_reshape, the rule now:
Extracts and honors the allowzero attribute from all matched Reshape nodes (shape_a, shape_b, and shape_c).
Resolves the input reshape shapes and computes the exact intermediate output shape of the matched MatMul using _compute_matmul_output_shape(...).
Resolves shape_c against matched_matmul_shape using allowzero, correctly preserving copying behavior (allowzero=0), literal zero behavior (allowzero=1), and -1 dimension inference.
Compares the resolved final output shape against the replacement's output shape before deciding whether to rewrite.
Added regression tests:
test_reshape_matmul_reshape_remain_when_non_identity_input_reshape_changes_output_shape: Covers Copilot's counterexample (A=[0, 0, 4], shape_a=[0, 0, 4], B=[1, 4, 5], shape_b=[1, 1, 4, 5], shape_c=[0, 0, 5], with allowzero=0). The final shape resolves to [1, 0, 5], which differs from the replacement's [0, 0, 5], so the rewrite is safely declined (count == 0).
test_reshape_matmul_reshape_replace_when_non_identity_input_reshape_allowzero_matches_shape: With allowzero=1, shape_c resolves to [0, 0, 5], matching the replacement, and the rewrite is accepted (count == 1).
test_reshape_matmul_reshape_replace_when_final_reshape_has_inferred_dimension: Verifies that inferred -1 dimensions in shape_c are correctly resolved and accepted when the shapes match.
All tests pass, including all 332 common rewriter-rule unit tests, and formatting has been verified with ruff.
Ruff check and format check pass for both changed files.
The remaining preview CI failures are reproducible on the unchanged base commit 06f8fc5: its Ubuntu ONNX-weekly check failed on both Oct 16 and Oct 19, and its Ubuntu/Windows torch-nightly checks failed on both dates. Codecov identifies them as pre-existing onnx_export optional-type cases, outside this MatMul rewrite.
The outstanding CodeRabbit thread has been addressed by the shape-resolution and allowzero regression coverage in 7fcb5f9 and is now resolved.
The reason will be displayed to describe this comment to others. Learn more.
Major: Newly accepted inferred final shapes expose transpose loss in the Gemm rewrite
Locations: onnxscript/rewriter/rules/common/_broadcast_to_matmul.py:314-328 and the direct consumer onnxscript/rewriter/rules/common/_gemm_to_matmul_add.py:8-22.
Resolving final reshape targets containing -1 broadens this shared condition beyond the MatMul rules. The Gemm rule also uses check_if_not_need_reshape, but its pattern does not constrain transA/transB, and its replacement emits a plain MatMul followed by Add, discarding any transpose.
The following valid opset-17 model now rewrites incorrectly:
<ir_version: 7, opset_import: [ "" : 17]>
agraph (float[2, 2] X, float[2, 2] Y, float[2, 2] C)
=> (float[2, 2] O)
{
sa = Constant<value: tensor = int64[2] {2, 2}>()
A = Reshape(X, sa)
M = Gemm<alpha=1.0, beta=1.0, transB=1>(A, Y, C)
sc = Constant<value: tensor = int64[2] {2, -1}>()
O = Reshape(M, sc)
}
With float32 X=eye(2), Y=[[1,2],[3,4]], and C=zeros((2,2)):
Original output: [[1,3],[2,4]].
Rewritten output: [[1,2],[3,4]].
Confirmed using both ONNX Runtime CPU with graph optimizations disabled and ONNX ReferenceEvaluator. Direct application of gemm_to_matmul_add_rule at 7fcb5f9 returns rewrite count 1 and changes the values. Isolated execution of the previous d3aa363f snapshot returns count 0 and preserves the original values. All original and rewritten models pass basic and full ONNX checking.
The missing transpose handling predates this commit, but this revision newly accepts this valid inferred-target graph, so the widened acceptance introduces an additional failing case. This is separate from the shared-shape/allowzero finding already reported by Copilot.
Suggested minimal fix: add a Gemm-specific guard that rejects nonzero transA/transB before using the shared reshape check, or preserve those transposes in the replacement. Add this inferred-final-shape regression with a nonsymmetric operand and assert numerical equivalence.
Scope: this repro directly applies the Gemm rewrite rule; it is not a full default optimization-pipeline run.
Addressed the Gemm transpose regression in a4cf747b.
Added a Gemm-specific guard that rejects transA != 0 or transB != 0 before calling the shared reshape condition, because the replacement emits plain MatMul + Add.
Added the supplied transB=1 / final [2, -1] model with nonsymmetric Y.
The regression asserts the rule returns 0 and uses ONNX ReferenceEvaluator to verify the output remains [[1, 3], [2, 4]].
Validation: Gemm suite 13 passed; common rules suite 334 passed, 1 skipped, 48 subtests passed; Ruff check/format pass.
The shared guard now accepts inferred final shapes but checks only the replacement MatMul, not the following Add. With input_a=[2] reshaped to [1,2], input_b=[2,3], input_c=[1,3], shape_c=[-1], and alpha=beta=1 without transposes, the original Gemm+Reshape outputs [3]. This guard accepts the rewrite, but Add broadcasts the replacement MatMul result back to [1,3], changing the output rank. Validate that the bias broadcasts to the replacement MatMul shape without expanding it, decline the rewrite when this cannot be established, and add this regression.
Added a Gemm-specific guard for transA/transB and a regression using the inferred final shape case you shared. The rule now declines that rewrite and preserves the expected output.
Validated with:
Gemm tests: 13 passed
Common rewriter tests: 334 passed, 1 skipped, 48 subtests passed
This confirms the existing concern in Copilot review 5578993581 with executable evidence; it is not a separate newly discovered issue. The previous shared-shape allowzero and Gemm transpose findings are resolved at this commit: their counterexamples now decline rewriting and preserve the original results.
The remaining guard verifies the replacement MatMul shape but does not account for the following Add with input_c. That Add can expand the output rank after the final Reshape has been removed.
Valid opset-17 reproducer:
<ir_version: 7, opset_import: [ "" : 17]>
agraph (float[2] X, float[2,3] Y, float[1,3] C)
=> (float[3] O)
{
sa = Constant<value: tensor = int64[2] {1,2}>()
A = Reshape(X,sa)
M = Gemm<alpha=1.0,beta=1.0>(A,Y,C)
sc = Constant<value: tensor = int64[1] {-1}>()
O = Reshape(M,sc)
}
Use float32 X=[1,2], Y=[[1,2,3],[4,5,6]], and C=[[10,20,30]].
Observed with both ONNX Runtime CPU (graph optimizations disabled) and ONNX ReferenceEvaluator:
Current rewrite count: 1.
Original output: [19,32,45], shape (3,).
Rewritten output: [[19,32,45]], shape (1,3).
Adding downstream Shape(O) changes its value from [3] to [1,3], making the semantic difference directly observable.
The raw replacement MatMul produces [3], which passes the shared guard, but adding C with shape [1,3] broadcasts the result back to [1,3]. A vector bias C with shape [3] is a passing control: rewriting preserves output shape [3].
Suggested fix: establish the complete replacement Add(MatMul(X,Y),C) output shape, not only MatMul shape. Reject the rewrite if bias broadcasting expands the expected final output shape or cannot be established statically, or retain the final Reshape. Add the rank-expanding bias regression above.
The missing bias check predates the latest transpose fix; accepting inferred final shapes in the earlier refactor exposes this additional failing case. Isolated d3aa363 source declines the reproducer (count 0).
Validation: 42 targeted broadcast/Gemm tests passed; Ruff lint and formatting checks passed. This report covers direct rule execution, not the full optimizer pipeline, GPU execution, or the full test suite.
1. Guard against bias broadcasting shape/rank expansion:
Updated check_if_not_need_reshape_for_gemm in _gemm_to_matmul_add.py to establish the complete replacement output shape for Add(MatMul(input_a, input_b), input_c) rather than checking MatMul output shape alone.
Rejects rewriting if input_c shape is not statically known or contains symbolic dimensions.
Computes multidirectional broadcast shape of MatMul(input_a, input_b) with input_c. If broadcasting fails or expands/alters the expected final output shape (replacement_add_shape != matmul_output_shape), the rewrite is safely declined.
Reproducer test (test_reshape_gemm_reshape_remain_when_bias_broadcasting_expands_output_rank): Confirms that for X: [2], Y: [2, 3], and C: [1, 3] with final target shape [3], the rewrite is declined (count == 0), preserving the runtime shape (3,) across both ReferenceEvaluator and ONNX Runtime CPU (with graph optimizations disabled).
Passing control test (test_reshape_gemm_reshape_replace_when_bias_vector_matches_final_shape): Confirms that when vector bias C has shape [3], the rewrite is accepted (count == 1) and preserves numerical accuracy and shape (3,).
Symbolic bias test (test_reshape_gemm_reshape_remain_when_bias_shape_has_symbolic_dim): Confirms that bias tensors with symbolic dimensions decline rewriting safely.
3. Validation:
All 45 targeted Gemm and broadcast-to-matmul tests passed (_gemm_to_matmul_add_test.py and _broadcast_to_matmul_test.py).
All 338 common rewriter tests passed (337 passed, 1 pre-existing skip).
Ruff check, Ruff format check, and Mypy static type checks passed cleanly.
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
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.
Ti-Tai Wang (titaiwangms)
left a comment
github-project-automation
Bot
moved this from Todo
to In Progress
in ONNX Script Review Board
Summary
In
_broadcast_to_matmul.py,check_if_not_need_reshapeinspected reversed input dimensions with:This introduced two issues:
idx == 0):At
idx == 0,dim_from_aanddim_from_bare the contracting dimensionsMatMul,dim_from_a == 1whendim_from_b != 1erroneously permitted rewriting valid reshape graphs into invalidMatMulnodes with mismatched contracting dimensions that fail ONNX shape inference.idx > 0):At
idx > 0,dim_from_aanddim_from_bare batch dimensions. Two dimensions are broadcastable ifChanges
idx == 0(dim_from_a == dim_from_b).idx > 0(dim_from_a == dim_from_b or dim_from_a == 1 or dim_from_b == 1)._broadcast_to_matmul_test.py:test_reshape_matmul_reshape_replace_when_second_input_batch_dim_is_one: confirms redundant reshapes are removed when inputtest_reshape_matmul_reshape_remain_when_contracting_dimension_mismatches: element-count-preserving fixture confirming the rewrite is safely declined when contracting dimensions mismatch.test_reshape_matmul_reshape_remain_when_reshapes_permute_batch_dimensions: confirms the rewrite is safely declined when reshapes alter batch-axis layout.test_reshape_matmul_reshape_numerical_equivalence_when_second_input_batch_dim_is_one: end-to-end regression verifying numerical equivalence between original and optimized models with ONNX Runtime.Validation
_broadcast_to_matmul_test.pypass.ruff checkandruff format --checkpass with zero warnings or errors.