From 2fd35eb7262cd2654895b08855ae917b1127396f Mon Sep 17 00:00:00 2001
From: liuquanbo-byte <liuquanbo@bytedance.com>
Date: Thu, 27 Jun 2024 17:44:23 +0800
Subject: [PATCH] [tf-frontend] infer shape after convert where to static shape
 op (#377)

infer shape after convert where to static shape op
---
 .../tf-frontend/tf_mlir_ext/pipelines/customized_tf_to_mhlo.cc   | 1 +
 1 file changed, 1 insertion(+)

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
index 659486907..83eb659cd 100644
--- 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(