From 6ad4a657e06c80ed71736a93b95e7b643606de1d Mon Sep 17 00:00:00 2001 From: Andy Ly Date: Wed, 8 Jul 2020 13:20:31 -0700 Subject: [PATCH] Add shape inference pass to the beginning of the TPUBridgeV1Compat pipeline. It is possible for the TPU computation to be in a function. Running shape inference prior to V1 specific passes allow for more granular shapes/types to be propagated to functions (e.g. resource subtypes), resulting in possibly less runtime/dynamic shapes for the TPU program. PiperOrigin-RevId: 320252292 Change-Id: I431ff20c943a9542fab2a0076675d62b22280cb8 --- tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc index 7b8ae474941..1963931b497 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc @@ -107,6 +107,7 @@ void CreateTPUBridgePipeline(OpPassManager &pm) { } void CreateTPUBridgePipelineV1(OpPassManager &pm) { + pm.addPass(TF::CreateTFShapeInferencePass()); // For V1 compatibility, we process a module where the graph does not have // feeds and fetched. We extract first the TPU computation in a submodule, // where it'll be in a function with args and returned values, much more like