From 048b336d448d89075357def0eba0cdee321f8081 Mon Sep 17 00:00:00 2001 From: jax authors Date: Wed, 5 Nov 2025 14:25:18 -0800 Subject: [PATCH] [Mosaic] Determine matmul data format from LHS instead of RHS. PiperOrigin-RevId: 828628903 --- jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc b/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc index 7a034f202ef0..ccee6bd0ddb4 100644 --- a/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc +++ b/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc @@ -713,6 +713,7 @@ FailureOr 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.