From 482da6730846fef3c1e42a7c490df7d697e0a594 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Sat, 17 Oct 2026 03:34:30 +0530 Subject: [PATCH] Fix contracting dim match and batch broadcasting in broadcast_to_matmul 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. --- .../rules/common/_broadcast_to_matmul.py | 12 ++++-- .../rules/common/_broadcast_to_matmul_test.py | 40 +++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py b/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py index ad673a7dcd..fb5b3e514e 100644 --- a/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py +++ b/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py @@ -100,10 +100,14 @@ def check_if_not_need_reshape( reversed(input_b_shape_except_last_dim), ) ): - if dim_from_a not in {1, dim_from_b}: - logger.info("Original shape is not broadcastable.") - return False - elif idx > 0: + if idx == 0: + if dim_from_a != dim_from_b: + logger.info("Original contracting dimensions do not match.") + return False + else: + if dim_from_a != dim_from_b and dim_from_a != 1 and dim_from_b != 1: + logger.info("Original shape is not broadcastable.") + return False broadcast_matmul_output_shape = [ max(dim_from_a, dim_from_b), # type: ignore[type-var] *broadcast_matmul_output_shape, diff --git a/onnxscript/rewriter/rules/common/_broadcast_to_matmul_test.py b/onnxscript/rewriter/rules/common/_broadcast_to_matmul_test.py index 8fcafb84a8..b16a98d793 100644 --- a/onnxscript/rewriter/rules/common/_broadcast_to_matmul_test.py +++ b/onnxscript/rewriter/rules/common/_broadcast_to_matmul_test.py @@ -370,6 +370,46 @@ def test_reshape_matmul_reshape_remain_when_output_is_not_matmul_broadcasted( self.assertEqual(count, 0) self.assertEqual(len(model.graph), 7) + def test_reshape_matmul_reshape_replace_when_second_input_batch_dim_is_one(self): + model_proto = onnx.parser.parse_model( + """ + + agraph (float[2, 3, 4] input_x, float[1, 4, 5] input_y) => (float[2, 3, 5] output) + { + shape_a = Constant() + reshape_x = Reshape (input_x, shape_a) + shape_b = Constant() + reshape_y = Reshape (input_y, shape_b) + matmul = MatMul (reshape_x, reshape_y) + shape_c = Constant() + output = Reshape (matmul, shape_c) + } + """ + ) + model = ir.serde.deserialize_model(model_proto) + count = _broadcast_to_matmul.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertEqual(len(model.graph), 4) + + def test_reshape_matmul_reshape_remain_when_contracting_dimension_mismatches(self): + model_proto = onnx.parser.parse_model( + """ + + agraph (float[3, 1] input_x, float[4, 5] input_y) => (float[3, 5] output) + { + shape_a = Constant() + reshape_x = Reshape (input_x, shape_a) + matmul = MatMul (reshape_x, input_y) + shape_c = Constant() + output = Reshape (matmul, shape_c) + } + """ + ) + model = ir.serde.deserialize_model(model_proto) + count = _broadcast_to_matmul.rules.apply_to_model(model) + self.assertEqual(count, 0) + self.assertEqual(len(model.graph), 5) + class OneReshapeMatMulReshapeTest(unittest.TestCase): def test_reshape_matmul_reshape_replace_when_nd_inputs_are_broadcastable(self):