From fac30b7a874baf4f2a326449b408e3a6f55c6c80 Mon Sep 17 00:00:00 2001 From: Feng Liu Date: Wed, 6 May 2020 17:44:29 -0700 Subject: [PATCH] Run shape inference before any new non-tf ops are introduced PiperOrigin-RevId: 310263391 Change-Id: Iff15999065f5ffa1c331646ae53c2f093c05b994 --- tensorflow/compiler/mlir/lite/tf_tfl_passes.cc | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc index 5eefa821c6b..6ab16141626 100644 --- a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc +++ b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc @@ -73,16 +73,17 @@ void AddTFToTFLConversionPasses(const mlir::TFL::PassConfig& pass_config, pass_manager->addPass(mlir::TFControlFlow::CreateRaiseTFControlFlowPass()); } + if (pass_config.shape_inference) { + pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass()); + } + // Keep this pass after the shape inference pass, which couldn't do shape + // inference for non-tf ops. if (!pass_config.quant_specs.serialized_quant_stats.empty()) { pass_manager->addPass( mlir::quant::CreateImportQuantStatsPassForTFControlDialect( pass_config.quant_specs.serialized_quant_stats)); } - if (pass_config.shape_inference) { - pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass()); - } - // The conversion pipeline has to follow the following orders: // 1) Saved model related optimization like decompose resource ops // 2) Convert composite functions like lstm/rnns, along with proper function