[Bug][Relax] CombineParallelMatmul fails valid branches with broadcast-compatible biases
The pass checks that bias tensors have the same rank, but does not check non-concat dimensions. It rewrites two valid broadcast additions into an invalid bias concat.
matmul(x, w1): [2, 4] + b1: [1, 4] -> [2, 4]
matmul(x, w2): [2, 5] + b2: [2, 5] -> [2, 5]
Both original additions are valid. The pass constructs concat((b1, b2), axis=1), which is invalid because the row dimensions are 1 and 2.
### Environment
OS: Linux x86_64
Target: llvm
Relax VM exec_mode: bytecode
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05
### Steps to reproduce
from tvm import relax
builder = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType((2, 3), "float32"))
w1 = relax.Var("w1", relax.TensorType((3, 4), "float32"))
w2 = relax.Var("w2", relax.TensorType((3, 5), "float32"))
b1 = relax.Var("b1", relax.TensorType((1, 4), "float32"))
b2 = relax.Var("b2", relax.TensorType((2, 5), "float32"))
with builder.function("main", [x, w1, w2, b1, b2]):
with builder.dataflow():
mm1 = builder.emit(relax.op.matmul(x, w1))
mm2 = builder.emit(relax.op.matmul(x, w2))
y1 = builder.emit(relax.op.add(mm1, b1))
y2 = builder.emit(relax.op.add(mm2, b2))
output = builder.emit_output((y1, y2))
builder.emit_func_output(output)
mod = relax.transform.Normalize()(builder.finalize())
relax.transform.CombineParallelMatmul()(mod)
Observed error:
ValueError: Concat expects the input tensors to have the same shape on every
dimension except the one indicated by the input axis. However, the input
contains tensors whose shapes on dimension 0 is T.int64(1) and T.int64(2)
Changing b2 to shape [1, 5] makes the concat valid; the original and transformed LLVM bytecode VM outputs then match exactly.
### Suspected cause
src/relax/transform/combine_parallel_matmul.cc records only a bias rank and later unconditionally builds:
auto concat_bias = concat(Tuple(bias), bias_dim - 1);
### Suggested fix
Validate compatibility of every bias shape on non-concat axes before applying the fused-bias rewrite. If that cannot be proven, combine only the matmuls and leave the individual bias additions in place.
### Triage
- needs-triage
- type: bug
- relax
```markdown
[Bug][Relax] CombineParallelMatmul fails valid branches with broadcast-compatible biases
Expected behavior
relax.transform.CombineParallelMatmul()should combine bias-add branches only when their biases can be concatenated along the output-channel axis. Otherwise it should leave the original valid branchesunchanged.
Actual behavior
The pass checks that bias tensors have the same rank, but does not check non-concat dimensions. It rewrites two valid broadcast additions into an invalid bias concat.