diff --git a/jaxlib/mosaic/dialect/tpu/tpu.td b/jaxlib/mosaic/dialect/tpu/tpu.td
index f5c27974385b..f2d57e86d20f 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu.td
+++ b/jaxlib/mosaic/dialect/tpu/tpu.td
@@ -1479,7 +1479,6 @@ def TPU_LogBufferOp : TPU_Op<"log_buffer"> {
 }
 
 
-
 def LogicalToPhysicalDeviceIdPass : Pass<"logical-to-physical-device-id", "::mlir::func::FuncOp"> {
   let dependentDialects = [
     "::mlir::func::FuncDialect",
@@ -1565,18 +1564,5 @@ def ApplyVectorLayoutPass : Pass<"tpu-apply-vector-layout", "::mlir::func::FuncO
   ];
 }
 
-def PreCanonicalizationOptimizationPass : Pass<"pre-canonicalization-optimization", "::mlir::func::FuncOp"> {
-  let summary = "Fold matmul rhs tranpose into the op before layout inference";
-  let constructor = "::mlir::tpu::createPreCanonicalizationOptimizationPass()";
-  let dependentDialects = [
-    "::mlir::vector::VectorDialect",
-    "::mlir::tpu::TPUDialect",
-  ];
-  let options = [
-    Option<"hardware_generation", "hardware-generation", "int", /*default=*/"6", "">,
-    Option<"lane_count", "lane-count", "int", /*default=*/"128", "">,
-    Option<"sublane_count", "sublane-count", "int", /*default=*/"8", "">,
-  ];
-}
 
 #endif  // TPU_ATTRS
diff --git a/jaxlib/mosaic/dialect/tpu/tpu_dialect.h b/jaxlib/mosaic/dialect/tpu/tpu_dialect.h
index 9ba415cef3f0..477fb04456d2 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu_dialect.h
+++ b/jaxlib/mosaic/dialect/tpu/tpu_dialect.h
@@ -90,11 +90,6 @@ std::unique_ptr<OperationPass<func::FuncOp>> createRelayoutInsertionPass(
 std::unique_ptr<OperationPass<func::FuncOp>> createApplyVectorLayoutPass(
     const ApplyVectorLayoutContext &ctx = ApplyVectorLayoutContext{});
 
-std::unique_ptr<OperationPass<func::FuncOp>>
-createPreCanonicalizationOptimizationPass(
-    int hardware_generation = -1,
-    std::array<int64_t, 2> target_shape = {8, 128});
-
 std::unique_ptr<OperationPass<func::FuncOp>>
 createLogicalToPhysicalDeviceIdPass(int64_t total_devices);