| 613 | } |
| 614 | |
| 615 | Status RewriteAssociatedFunction( |
| 616 | Graph* graph, Node* node, FunctionLibraryDefinition* fld, |
| 617 | const AssociatedFunctionInfo& associated_function, |
| 618 | const string& rewritten_function_name) { |
| 619 | switch (associated_function.type()) { |
| 620 | case AssociatedFunctionInfo::kFunctionCallNode: { |
| 621 | // Change this node to call the new function. |
| 622 | NodeDebugInfo debug_info(*node); |
| 623 | NodeDefBuilder builder(node->name(), rewritten_function_name, fld, |
| 624 | &debug_info); |
| 625 | for (auto attr : node->attrs()) { |
| 626 | builder.Attr(attr.first, attr.second); |
| 627 | } |
| 628 | for (int i = 0; i < node->num_inputs(); i++) { |
| 629 | Node* input_node; |
| 630 | TF_RETURN_IF_ERROR(node->input_node(i, &input_node)); |
| 631 | builder.Input(input_node->name(), i, node->input_type(i)); |
| 632 | } |
| 633 | builder.Device(node->assigned_device_name().empty() |
| 634 | ? node->requested_device() |
| 635 | : node->assigned_device_name()); |
| 636 | NodeDef node_def; |
| 637 | TF_RETURN_IF_ERROR(builder.Finalize(&node_def)); |
| 638 | Status s; |
| 639 | Node* new_node = graph->AddNode(node_def, &s); |
| 640 | TF_RETURN_IF_ERROR(s); |
| 641 | for (auto edge : node->in_edges()) { |
| 642 | graph->AddEdge(edge->src(), edge->src_output(), new_node, |
| 643 | edge->dst_input()); |
| 644 | } |
| 645 | for (auto edge : node->out_edges()) { |
| 646 | graph->AddEdge(new_node, edge->src_output(), edge->dst(), |
| 647 | edge->dst_input()); |
| 648 | } |
| 649 | graph->RemoveNode(node); |
| 650 | break; |
| 651 | } |
| 652 | case AssociatedFunctionInfo::kSymbolicGradient: { |
| 653 | NameAttrList func; |
| 654 | TF_RETURN_IF_ERROR(GetNodeAttr( |
| 655 | node->attrs(), FunctionLibraryDefinition::kFuncAttr, &func)); |
| 656 | GradientDef gradient_def; |
| 657 | gradient_def.set_function_name(func.name()); |
| 658 | gradient_def.set_gradient_func(rewritten_function_name); |
| 659 | string original_grad_func = fld->FindGradient(func.name()); |
| 660 | if (original_grad_func.empty()) { |
| 661 | TF_RETURN_IF_ERROR(fld->AddGradientDef(gradient_def)); |
| 662 | } else if (original_grad_func != rewritten_function_name) { |
| 663 | TF_RETURN_IF_ERROR(fld->ReplaceGradient(gradient_def)); |
| 664 | } |
| 665 | break; |
| 666 | } |
| 667 | case AssociatedFunctionInfo::kFunctionAttr: { |
| 668 | // Change function attr to rewritten functions. |
| 669 | NameAttrList func; |
| 670 | TF_RETURN_IF_ERROR( |
| 671 | GetNodeAttr(node->attrs(), associated_function.attr_name(), &func)); |
| 672 | node->ClearAttr(associated_function.attr_name()); |
no test coverage detected