| 502 | } // namespace |
| 503 | |
| 504 | Status MakeFunctionDef(const GrapplerFunctionItem& item, |
| 505 | const FunctionLibraryDefinition& flib, |
| 506 | FunctionDef* func) { |
| 507 | func->mutable_signature()->set_name(item.id); |
| 508 | func->mutable_signature()->set_description(item.description()); |
| 509 | func->mutable_signature()->set_is_stateful(item.is_stateful()); |
| 510 | |
| 511 | MakeFunctionDefHelper helper; |
| 512 | TF_RETURN_IF_ERROR(helper.Initialize(item, flib)); |
| 513 | |
| 514 | // Mapping from the '_Retval' node name to the output tensor. |
| 515 | absl::flat_hash_map<absl::string_view, string> output_tensors; |
| 516 | for (const NodeDef& func_body_node : item.function_body().node()) { |
| 517 | if (!helper.IsOutputNode(func_body_node)) continue; |
| 518 | if (func_body_node.input_size() != 1) { |
| 519 | return errors::Internal("_Retval node must have single input: ", |
| 520 | SummarizeNodeDef(func_body_node)); |
| 521 | } |
| 522 | output_tensors.emplace(func_body_node.name(), func_body_node.input(0)); |
| 523 | } |
| 524 | |
| 525 | for (const InputArgInstantiation& input_arg : item.inputs()) { |
| 526 | OpDef::ArgDef arg_def; |
| 527 | arg_def.set_name(input_arg.node_name); |
| 528 | arg_def.set_type(input_arg.data_type); |
| 529 | arg_def.set_is_ref(IsRefType(input_arg.data_type)); |
| 530 | *func->mutable_signature()->add_input_arg() = arg_def; |
| 531 | } |
| 532 | |
| 533 | // Add function output arguments. |
| 534 | for (const OutputArgInstantiation& output_arg : item.outputs()) { |
| 535 | const string output_name = |
| 536 | absl::StrReplaceAll(output_arg.node_name, {{"_RetVal", ""}}); |
| 537 | |
| 538 | OpDef::ArgDef arg_def; |
| 539 | arg_def.set_name(output_name); |
| 540 | arg_def.set_type(output_arg.data_type); |
| 541 | arg_def.set_is_ref(IsRefType(output_arg.data_type)); |
| 542 | *func->mutable_signature()->add_output_arg() = arg_def; |
| 543 | |
| 544 | auto it = output_tensors.find(output_arg.node_name); |
| 545 | if (it == output_tensors.end()) { |
| 546 | return errors::Internal( |
| 547 | "Can't find an output tensor for the output node: ", |
| 548 | output_arg.node_name); |
| 549 | } |
| 550 | |
| 551 | TF_RETURN_IF_ERROR(helper.AsFunctionDefInput( |
| 552 | it->second, &(*func->mutable_ret())[output_name])); |
| 553 | } |
| 554 | |
| 555 | // Add function control outputs. |
| 556 | for (const ControlOutput& control_out : item.control_outputs()) { |
| 557 | func->mutable_control_ret()->insert( |
| 558 | {control_out.output_name, control_out.node_name}); |
| 559 | *func->mutable_signature()->add_control_output() = control_out.output_name; |
| 560 | } |
| 561 | |