Skip to content

[Fix] Batch MatMul Op#3413

Open
veblush wants to merge 1 commit intotensorflow:mainfrom
veblush:matmul
Open

[Fix] Batch MatMul Op#3413
veblush wants to merge 1 commit intotensorflow:mainfrom
veblush:matmul

Conversation

@veblush
Copy link
Collaborator

@veblush veblush commented Feb 9, 2026

Batmul

  • Disabled optimization when adj_x=true.
  • Extracted logic into ReshapeToFlattenBatchDimsIfPossible for better maintainability.
  • Added two new tests.

BUG=#3412

@TFLM-bot TFLM-bot added the ci:ready Triggers the basic TFLM test suite. label Feb 9, 2026
@veblush veblush marked this pull request as ready for review February 9, 2026 18:41
@veblush veblush requested a review from a team as a code owner February 9, 2026 18:41
Copy link
Member

@ddavis-2015 ddavis-2015 left a comment

Choose a reason for hiding this comment

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

@veblush The CMSIS_NN optimized kernel also needs to be updated.

Comment on lines +353 to +364
// We can only flatten the dimensions if the physical layout in memory
// allows us to treat [Batch, Row, Col] as [Batch * Row, Col].
// This requires the 'Row' dimension to be contiguous with the 'Batch'
// dimension.
//
// If adj_x is true, the logical operation is on transposed matrices.
// The physical layout is [Batch, Row, Col] but logically we access [Batch,
// Col, Row]. Flattening Batch and Row dimensions (physically) results in
// [Batch*Row, Col]. This does not match the logical expectation of [Batch,
// Col, Row] because the columns of the second batch are not contiguous with
// the columns of the first batch in the logical transposed view. Therefore,
// we disable this optimization when adj_x is true.
Copy link
Member

Choose a reason for hiding this comment

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

Hmmm. This means the LiteRT code also needs to be corrected.

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

Labels

ci:ready Triggers the basic TFLM test suite.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants