[Mosaic] Fix canonicalize_extract op name.

PiperOrigin-RevId: 694236671
jax-ml · Nov 7, 2024 · 8b7bcad · 8b7bcad
1 parent 0b022aa
commit 8b7bcad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc b/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
@@ -528,7 +528,7 @@ const llvm::StringMap<canonicalize_rule_type> &rules() {
   static auto rules = new llvm::StringMap<canonicalize_rule_type>{
       {tpu::MatmulOp::getOperationName(), canonicalize_matmul},
       {vector::ContractionOp::getOperationName(), canonicalize_contraction},
-      {vector::ContractionOp::getOperationName(), canonicalize_extract},
+      {vector::ExtractOp::getOperationName(), canonicalize_extract},
       {vector::MultiDimReductionOp::getOperationName(),
        canonicalize_multi_dim_reduction},
       {arith::SelectOp::getOperationName(), canonicalize_select},