experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix shape inference call in compile op

Browse Source 
PiperOrigin-RevId: 337246090
Change-Id: I7053df6c4d7bb61244ad0b0081904156bffa975a
This commit is contained in:
Yuanzhong Xu 2020-10-14 23:01:29 -07:00 committed by TensorFlower Gardener
parent 5452c25097
commit 19200d4ec5
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
View File

@ -514,8 +514,7 @@ Status TpuCompileOpKernelCommon::OptimizeGraph(
// Converts the GraphShapeInfo into the form needed by the constant-folding // Converts the GraphShapeInfo into the form needed by the constant-folding
// pass of the optimizer. // pass of the optimizer.
std::unordered_map<string, std::vector<PartialTensorShape>> shape_map; std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
TF_RETURN_IF_ERROR(RunShapeInferenceOnComputation( ConvertGraphShapeInfoToShapeMap(**graph, shape_info, &shape_map);
metadata, arg_shapes, graph->get(), flr, &shape_info));
optimizer_opts.shape_map = &shape_map; optimizer_opts.shape_map = &shape_map;
optimizer.Optimize(flr, flr->env(), flr->device(), graph, optimizer_opts); optimizer.Optimize(flr, flr->env(), flr->device(), graph, optimizer_opts);
} }