fuse_matmul_weight_bias: fix sub operand order and drop the const-on-the-left case

[LeSingh1]

Problem

fuse_matmul_weight_bias rewrites matmul + add/sub into linear. I swept all 32 combinations of (which matmul operand is const) × transpose_x × transpose_y × add/sub × (which add/sub operand the matmul output is), evaluating the program before and after the pass. 20 of the 32 fuse into a program that computes something different. There are two independent causes.

1. sub is treated as commutative

The docstring says:

%5 = add(x=%3, y=%4) # %4 is const. add(x=%4, y=%3) is equivalent
                     # sub is similar.

It is not similar — sub is not commutative. The code's entire sub handling is:

if add_op.op_type == "sub":
    bias = -bias

with no is_first_input test anywhere. So bias - x @ W silently becomes x @ W - bias:

@mb.program(input_specs=[mb.TensorSpec(shape=(2, 3))])
def prog(x):
    mm = mb.matmul(x=x, y=np.eye(3, dtype=np.float32))
    return mb.sub(x=np.array([100., 200., 300.], np.float32), y=mm)

With x = [[1,2,3],[4,5,6]]:

	value
before	[[99, 198, 297], [96, 195, 294]]
after	[[-99, -198, -297], [-96, -195, -294]]

Exactly negated. apply_pass_and_basic_check raises nothing — the shapes and output names are unchanged.

fuse_linear_bias, in the same file, already gets this right: it computes is_first_input and negates the weight when the fused op is the second operand.

2. Const as the left matmul operand puts the bias inside a transpose

w @ x is rewritten as transpose(linear(transpose(x), w)) with the bias inside the linear. But linear adds its bias along its own last axis, which the outer transpose then moves to axis −2 of the result, so the bias lands on the wrong axis:

mm  = mb.matmul(x=np.eye(3, dtype=np.float32), y=x)
out = mb.add(x=mm, y=np.array([100., 200., 300.], np.float32))

With x = np.arange(9).reshape(3, 3):

	value
before	[[100,201,302],[103,204,305],[106,207,308]]
after	[[100,101,102],[203,204,205],[306,307,308]]

The existing guard cannot catch this: d_out is taken from the weight, but in this branch the output's last dimension comes from the other operand, so bias.shape[0] != d_out only lets the fusion through when everything is square — exactly the case where the shape check cannot notice.

All 16 const_is_x=True combinations are wrong, plus the 4 sub-with-bias-first ones from cause 1.

common::fuse_matmul_weight_bias is in the default pipeline (pass_pipeline.py).

Fix

• For sub, test which operand the matmul output is. x @ w - bias stays linear(x, w, -bias); bias - x @ w becomes linear(x, -w, bias), mirroring fuse_linear_bias.
• Bail out when the const is the left matmul operand. The bias genuinely cannot be expressed inside the linear there, and emitting it as a separate add after the transpose would turn 2 ops into 4 — not a fusion at all. This removes an "optimization" that was wrong in every one of its 16 configurations.

After the change the same 32-case sweep fuses 16 and all 16 are numerically identical to the original.

Tests

In TestFuseMatmulWeightBias (passes/tests/test_optimize_linear_passes.py):

• test_fuse_matmul_weight_bias_sub_with_matmul_second[transpose_x, transpose_y] — 4 cases; fuses and then compares the evaluated result against the original. Fails on main.
• test_no_fuse_matmul_weight_bias_const_on_the_left[op_type, bias_first] — 4 cases; the program must be left alone. Fails on main.

The tests evaluate the programs with a small interpreter over the five op types involved (matmul, linear, transpose, add, sub) so the numbers are checked without needing a backend.

The existing test_fuse_matmul_weight_bias is untouched and its fusion assertions still pass. I ran test_optimize_linear_passes.py in full plus test_passes.py::TestFuseMatmulWeightBias before and after; the pre-existing failure set is identical, 56 either way (this environment cannot load CoreML.framework, so assert_model_is_valid fails there regardless).