Fold Gemm beta into the fused BatchNorm bias

[Mohammed Alkindi (MohammedAlkindi)]

FuseBatchNormIntoGemm folds BatchNormalization(Gemm(x)) into one Gemm, but the bias math assumes beta is 1. Gemm computes Y = alpha*A'B' + beta*C, so with beta != 1 the fused bias is wrong twice over: beta is never applied to original_bias, and the original attribute is forwarded onto the new node, where it scales the already-folded bias a second time.

The rule is in _DEFAULT_REWRITE_RULES, so this fires from a plain rewrite(model) with no configuration, and it fails silently, producing a wrong tensor with no exception.

On Gemm<beta=2.0>(X, W, B) followed by BatchNorm, comparing the original model against the rewritten one through onnxruntime:

Mismatched elements: 7 / 64 (10.9%)
Max relative difference among violations: 22.966587

Conv and ConvTranspose have no beta, so they take the 1.0 default and are unaffected.

Two parameterized cases added (beta=0.5, beta=2.0), failing before and passing after. Reverting either half of the fix on its own, the bias scaling or the attribute drop, makes them fail again, so both are load-bearing.

onnxscript/rewriter/ is 508 passed, 2 skipped, with gqa_test.py and models/ excluded because torch is not installed here. ruff check and ruff format --check clean.

[Copilot]

🟢 Approval recommended

The fix is narrowly scoped, removes a clear double-application bug, and is backed by new regression tests that validate both numerical correctness and model validity.

Pull request overview

Fixes a correctness bug in the default BatchNorm fusion rule when fusing BatchNormalization(Gemm(...)) for Gemm nodes with non-default beta (bias scaling) values. This prevents silently incorrect rewritten models by folding Gemm’s bias scaling into the fused bias and ensuring the attribute isn’t applied twice.

Changes:

• Fold Gemm’s beta attribute into the fused bias computation during BatchNorm fusion.
• Drop the beta attribute from the re-emitted fused Gemm node to avoid double-scaling.
• Add parameterized regression tests covering beta=0.5 and beta=2.0.

File summaries

File	Description
onnxscript/rewriter/rules/common/_fuse_batchnorm.py	Corrects Gemm+BatchNorm fusion by folding Gemm’s beta scaling into the fused bias and removing the forwarded beta attribute.
onnxscript/rewriter/rules/common/_fuse_batchnorm_test.py	Adds parameterized tests to validate numerical equivalence for Gemm beta != 1.0 after fusion.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.68%. Comparing base (995ce98) to head (47d71e9).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3027      +/-   ##
==========================================
+ Coverage   72.67%   72.68%   +0.01%     
==========================================
  Files         265      265              
  Lines       32266    32281      +15     
  Branches     3052     3052              
==========================================
+ Hits        23448    23463      +15     
  Misses       7782     7782              
  Partials     1036     1036              

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