| 2338 | } |
| 2339 | |
| 2340 | bool ConstantFolding::SimplifyPack(GraphDef* optimized_graph, NodeDef* node) { |
| 2341 | if (!(IsPack(*node) && NumNonControlInputs(*node) == 1 && |
| 2342 | !OptimizedNodeExists(*node, "_const_axis"))) { |
| 2343 | return false; |
| 2344 | } |
| 2345 | // Create constant axis node. |
| 2346 | Tensor axis_t(DT_INT32, TensorShape({})); |
| 2347 | NodeDef* axis_node = optimized_graph->add_node(); |
| 2348 | axis_node->set_name(OptimizedNodeName(*node, "_const_axis")); |
| 2349 | const int axis = |
| 2350 | node->attr().count("axis") == 0 ? 0 : node->attr().at("axis").i(); |
| 2351 | if (!SetTensorValue(DT_INT32, axis, &axis_t).ok() || |
| 2352 | !CreateNodeDef(axis_node->name(), TensorValue(&axis_t), axis_node).ok()) { |
| 2353 | return false; |
| 2354 | } |
| 2355 | // Add a control dependency to make sure axis_node is in the right frame. |
| 2356 | const string ctrl_dep = ConstantFolding::AddControlDependency( |
| 2357 | node->input(0), optimized_graph, node_map_.get()); |
| 2358 | axis_node->add_input(ctrl_dep); |
| 2359 | axis_node->set_device(node->device()); |
| 2360 | node->set_op("ExpandDims"); |
| 2361 | if (node->attr().count("axis") != 0) { |
| 2362 | node->mutable_attr()->erase("axis"); |
| 2363 | } |
| 2364 | if (node->attr().count("N") != 0) { |
| 2365 | node->mutable_attr()->erase("N"); |
| 2366 | } |
| 2367 | (*node->mutable_attr())["Tdim"].set_type(DT_INT32); |
| 2368 | node->add_input(axis_node->name()); |
| 2369 | if (node->input_size() > 2) { |
| 2370 | node->mutable_input()->SwapElements(1, node->input_size() - 1); |
| 2371 | } |
| 2372 | return true; |
| 2373 | } |
| 2374 | |
| 2375 | bool ConstantFolding::MoveConstantsPastEnter(GraphDef* optimized_graph, |
| 2376 | NodeDef* node) { |
nothing calls this directly
no test coverage detected