OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-07 12:20:24 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Add shape inference pass to the beginning of the TPUBridgeV1Compat pipeline.

Browse Source 
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
This commit is contained in:
Andy Ly 2020-07-08 13:20:31 -07:00 committed by TensorFlower Gardener
parent 26e2f2e285
commit 6ad4a657e0
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc
View File

@ -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