Skip to content

Re-enable 3D x 2D BatchMatMul dimension collapse for adj_x=true#3599

Open
veblush wants to merge 2 commits into
tensorflow:mainfrom
veblush:bmat
Open

Re-enable 3D x 2D BatchMatMul dimension collapse for adj_x=true#3599
veblush wants to merge 2 commits into
tensorflow:mainfrom
veblush:bmat

Conversation

@veblush

@veblush veblush commented Jun 16, 2026

Copy link
Copy Markdown
Collaborator

Re-enables the outer dimension collapse optimization ([B, C, K] -> [B*C, K]) for 3D x 2D BatchMatMul when adj_x = true, removing the workaround from PR #3413.

By executing dimension collapse after shape assignment, lhs_shape is in the correct canonical form to safely collapse [Batch, Col] into [Batch * Col] across both ReferenceOps and CMSIS-NN. Added unit tests to verify and protect this optimization.

BUG=#3412

@veblush veblush requested a review from a team as a code owner June 16, 2026 23:23
@veblush veblush added the ci:full Triggers the comprehensive cross-platform test suite. label Jun 16, 2026
@ddavis-2015 ddavis-2015 added ci:full Triggers the comprehensive cross-platform test suite. and removed ci:full Triggers the comprehensive cross-platform test suite. labels Jun 17, 2026

@ddavis-2015 ddavis-2015 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good to me.

@veblush veblush requested a review from suleshahid June 18, 2026 16:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:full Triggers the comprehensive cross-platform test suite.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants