Converts any sub-graphs that can be resolved into constant expressions into single Const ops.
| 211 | // Converts any sub-graphs that can be resolved into constant expressions into |
| 212 | // single Const ops. |
| 213 | Status FoldConstants(const GraphDef& input_graph_def, |
| 214 | const TransformFuncContext& context, |
| 215 | GraphDef* output_graph_def) { |
| 216 | Graph input_graph(OpRegistry::Global()); |
| 217 | TF_RETURN_IF_ERROR(input_graph.AddFunctionLibrary(input_graph_def.library())); |
| 218 | |
| 219 | ShapeRefiner shape_refiner(input_graph.versions(), input_graph.op_registry()); |
| 220 | shape_refiner.set_require_shape_inference_fns(false); |
| 221 | shape_refiner.set_disable_constant_propagation(false); |
| 222 | shape_refiner.set_function_library_for_shape_inference( |
| 223 | &input_graph.flib_def()); |
| 224 | |
| 225 | bool clear_output_shapes; |
| 226 | TF_RETURN_IF_ERROR(context.GetOneBoolParameter("clear_output_shapes", true, |
| 227 | &clear_output_shapes)); |
| 228 | if (clear_output_shapes) { |
| 229 | // Some older GraphDefs have saved _output_shapes attributes which are out |
| 230 | // of date and cause import errors, so clean them up first. |
| 231 | GraphDef cleaned_graph_def; |
| 232 | RemoveAttributes(input_graph_def, {"_output_shapes"}, &cleaned_graph_def); |
| 233 | |
| 234 | TF_RETURN_IF_ERROR( |
| 235 | ImportGraphDef({}, cleaned_graph_def, &input_graph, &shape_refiner)); |
| 236 | } else { |
| 237 | TF_RETURN_IF_ERROR( |
| 238 | ImportGraphDef({}, input_graph_def, &input_graph, &shape_refiner)); |
| 239 | } |
| 240 | |
| 241 | // Sorted array of input names as lookup table. |
| 242 | std::vector<TensorId> input_names; |
| 243 | input_names.reserve(context.input_names.size()); |
| 244 | std::transform(context.input_names.begin(), context.input_names.end(), |
| 245 | std::back_inserter(input_names), |
| 246 | [](const string& name) { return ParseTensorName(name); }); |
| 247 | |
| 248 | const auto compare = [](TensorId lhs, TensorId rhs) { |
| 249 | return lhs.first < rhs.first; |
| 250 | }; |
| 251 | |
| 252 | std::sort(input_names.begin(), input_names.end(), compare); |
| 253 | |
| 254 | // Set statically inferred shapes. |
| 255 | std::unordered_map<string, std::vector<PartialTensorShape>> shape_map; |
| 256 | for (const Node* const node : input_graph.nodes()) { |
| 257 | auto ctx = shape_refiner.GetContext(node); |
| 258 | if (ctx == nullptr) { |
| 259 | continue; |
| 260 | } |
| 261 | |
| 262 | std::vector<PartialTensorShape>& partial_shapes = shape_map[node->name()]; |
| 263 | if (ctx->num_outputs() <= 0) continue; |
| 264 | partial_shapes.resize(ctx->num_outputs()); |
| 265 | |
| 266 | // Check all outputs. |
| 267 | for (const Edge* out_edge : node->out_edges()) { |
| 268 | if (out_edge->IsControlEdge()) continue; |
| 269 | |
| 270 | const int output_idx = out_edge->src_output(); |