Skip to content

[FIX][Relax] Skip ExpandMatmulOfSum when addition broadcasts - #20714

Open
tandede wants to merge 1 commit into
apache:mainfrom
tandede:fix/expand-matmul-of-sum-broadcast
Open

tandede wants to merge 1 commit into
apache:mainfrom
tandede:fix/expand-matmul-of-sum-broadcast

Conversation

@tandede

@tandede tandede commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

Fixes #20697.

ExpandMatmulOfSum currently distributes matmul(x, A + B) even when A and B are broadcast to form the sum. In that case, the independently computed matmuls can have different shapes, and adding them may produce a different result shape from the original expression.

Skip the rewrite unless both addends have the same inferred tensor type as the sum. This preserves the existing optimization for non-broadcast additions while avoiding the invalid transformation.

A regression test covers addends with shapes [1] and [1, 1, 1].

Tests:

  • tests/python/relax/test_transform_expand_matmul_of_sum.py
  • pre-commit checks for the modified files

@tandede
tandede force-pushed the fix/expand-matmul-of-sum-broadcast branch from 51912a1 to 1b7a4ee Compare October 15, 2026 02:06

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] ExpandMatmulOfSum silently miscompiles matmul(x, a+b) when addends broadcast

1 participant