| 296 | } |
| 297 | |
| 298 | Status ReplaceInputWithConst(const NodeDef& input_const, int input_index, |
| 299 | GrapplerFunctionItem* item) { |
| 300 | if (!IsConstant(input_const)) { |
| 301 | return errors::InvalidArgument("Input node is not a constant: ", |
| 302 | SummarizeNodeDef(input_const)); |
| 303 | } |
| 304 | if (input_index < 0 || input_index >= item->input_size()) { |
| 305 | return errors::InvalidArgument( |
| 306 | "Function input index is out of bound: index=", input_index, |
| 307 | " input_size=", item->input_size()); |
| 308 | } |
| 309 | |
| 310 | const InputArgInstantiation& input_arg = item->input(input_index); |
| 311 | |
| 312 | for (NodeDef& node : *item->graph.mutable_node()) { |
| 313 | // Replace '_Arg' node in the function body with a 'Const' node. |
| 314 | if (node.name() == input_arg.node_name) { |
| 315 | node = input_const; |
| 316 | node.set_name(input_arg.node_name); |
| 317 | node.clear_input(); |
| 318 | node.clear_device(); // device placement is defined by instantiating node |
| 319 | } |
| 320 | |
| 321 | // Update index in all inputs after the removed const input. |
| 322 | if (IsArg(node)) { |
| 323 | auto attrs = AttrSlice(node); |
| 324 | int index; |
| 325 | TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "index", &index)); |
| 326 | if (index >= input_index) { |
| 327 | (*node.mutable_attr())["index"].set_i(index - 1); |
| 328 | } |
| 329 | } |
| 330 | } |
| 331 | |
| 332 | item->input_args_.erase(item->input_args_.begin() + input_index); |
| 333 | |
| 334 | return Status::OK(); |
| 335 | } |
| 336 | |
| 337 | Status RemoveFunctionOutputs(const absl::flat_hash_set<int>& remove_outputs, |
| 338 | GrapplerFunctionItem* item, |