[Mosaic] Determine matmul data format from LHS instead of RHS.

Google-ML-Automation · Google-ML-Automation · commit 048b336d448d · 2025-11-14T08:07:03.000-08:00
PiperOrigin-RevId: 828628903
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc b/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
@@ -713,6 +713,7 @@ FailureOr<Value> canonicalize_matmul(const CanonicalizeContext &ctx,
     }
   }
 
+
   // Attempt to canonicalize matmul(x, transpose(y)) to a matmul with the
   // dimension numbers changed which will later be lowered into a more efficient
   // operation that fuses the transpose into the matmul.

Original file line number	Diff line number	Diff line change
`@@ -713,6 +713,7 @@ FailureOr<Value> canonicalize_matmul(const CanonicalizeContext &ctx,`
`713`	`713`	`}`
`714`	`714`	`}`
`715`	`715`
	`716`	`+`
`716`	`717`	`// Attempt to canonicalize matmul(x, transpose(y)) to a matmul with the`
`717`	`718`	`// dimension numbers changed which will later be lowered into a more efficient`
`718`	`719`	`// operation that fuses the transpose into the matmul.`