[tf-frontend] infer shape after convert where to static shape op (#377)

infer shape after convert where to static shape op
bytedance · Jun 27, 2024 · 2fd35eb · 2fd35eb
1 parent bed7740
commit 2fd35eb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/frontends/tf-frontend/tf_mlir_ext/pipelines/customized_tf_to_mhlo.cc b/frontends/tf-frontend/tf_mlir_ext/pipelines/customized_tf_to_mhlo.cc
@@ -187,6 +187,7 @@ struct CustomizedTfToMhloPipelinePass
       return;
     }
 
+    pm.addPass(mlir::TF::CreateTFShapeInferencePass());
     pm.addNestedPass<mlir::func::FuncOp>(
         mlir::tfext::createMhloLegalizeTfExtPass());
     pm.addPass(mlir::mhlo::createLegalizeTFPass(