Fold Gemm beta into the fused BatchNorm bias

onnxscript/rewriter/rules/common/_fuse_batchnorm.py

@@ -69,6 +69,14 @@ def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Valu
             self._scale_weights(weights, scale_factor, inbound_node.attributes)
         )
 
+        # Gemm optionally scales its bias input C by the "beta" attribute (default 1.0,
+        # https://onnx.ai/onnx/operators/onnx__Gemm.html#attributes); Conv/ConvTranspose have
+        # no such attribute. Fold that scaling into the fused bias now, and drop "beta" from
+        # the re-emitted node's attributes so it is not applied a second time there.
+        new_attributes = dict(inbound_node.attributes)
+        gemm_beta_attr = new_attributes.pop("beta", None)
+        gemm_beta = gemm_beta_attr.as_float() if gemm_beta_attr is not None else 1.0
+
         # Update bias
         if len(inbound_node.inputs) > 2:
             original_bias = inbound_node.inputs[2].const_value.numpy()
@@ -79,14 +87,14 @@ def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Valu
             # to avoid name collision on initializer creation when there are multiple patterns
             # sharing the same parent nodes.
             bias_name = inbound_node.inputs[1].name + "_bias"
-        fused_bias = ir.tensor((original_bias - input_mean) * scale_factor + beta)
+        fused_bias = ir.tensor((gemm_beta * original_bias - input_mean) * scale_factor + beta)
 
         return op.op(
             self.op_type,
             x,
             op.initializer(fused_weights, name=inbound_node.inputs[1].name),
             op.initializer(fused_bias, name=bias_name),
-            **inbound_node.attributes,
+            **new_attributes,
         )
 
     def check(self, context, x, inbound_out: ir.Value, batchnorm_out: ir.Value) -> MatchResult:

onnxscript/rewriter/rules/common/_fuse_batchnorm_test.py

@@ -396,6 +396,48 @@ def test_fuse_batchnorm_skips_shared_weight_initializers(self):
             ),
         )
 
+    @parameterized.parameterized.expand(
+        [
+            ("beta_half", 0.5),
+            ("beta_two", 2.0),
+        ]
+    )
+    def test_fuse_batchnorm_gemm_scales_bias_by_beta(self, _: str, beta_value: float):
+        """Gemm's beta scales input C, so it must be folded into the fused bias."""
+        model_proto = onnx.parser.parse_model(f"""
+            < ir_version: 7, opset_import: ["" : 17] >
+            test_model (float[N, 32] X) => (float [N, ?] Y)
+            <float[32, 64] W, float[64] B, float[64] gamma,
+             float[64] beta, float[64] input_mean, float[64] input_var>
+            {{
+                X1 = Gemm<beta={beta_value}>(X, W, B)
+                Y = BatchNormalization(X1, gamma, beta, input_mean, input_var)
+            }}
+        """)
+        model_proto.graph.initializer.extend(
+            [
+                onnx.numpy_helper.from_array(
+                    np.random.randn(32, 64).astype(np.float32), name="W"
+                ),
+                onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="B"),
+                *self._create_batchnorm_params(size=64),
+            ]
+        )
+
+        onnx.checker.check_model(model_proto, True)
+        model = ir.serde.deserialize_model(model_proto)
+
+        count = _fuse_batchnorm.rules.apply_to_model(model)
+        self.assertEqual(count, 1)
+        self.assertEqual(len(model.graph), 1)
+
+        testing.assert_numerically_equal(
+            model_proto, model, (np.random.rand(1, 32).astype(np.float32),)
+        )
+
+        output_model_proto = ir.serde.serialize_model(model)
+        onnx.checker.check_model(output_model_proto, True)
+
 
 if __name__ == "__main__":
     unittest.main()