| 584 | } |
| 585 | |
| 586 | std::vector<AssociatedFunctionInfo> GetAssociatedFunctions( |
| 587 | const Node& node, const FunctionLibraryDefinition* fld) { |
| 588 | std::vector<AssociatedFunctionInfo> results; |
| 589 | const string& op = node.type_string(); |
| 590 | if (fld->Contains(op)) { |
| 591 | // This is a function call node. |
| 592 | AttrValueMap attrs(node.attrs().begin(), node.attrs().end()); |
| 593 | results.emplace_back(AssociatedFunctionInfo::FunctionCall(op, attrs)); |
| 594 | } else if (node.type_string() == FunctionLibraryDefinition::kGradientOp) { |
| 595 | // This is a SymbolicGradient op. |
| 596 | AttrValueMap attrs(node.attrs().begin(), node.attrs().end()); |
| 597 | results.emplace_back(AssociatedFunctionInfo::SymbolicGradient(op, attrs)); |
| 598 | } else if (node.type_string() == "XlaHostCompute") { |
| 599 | // XlaHostCompute has "shape_inference_graph" func attr, but that's not |
| 600 | // related to graph execution. |
| 601 | } else { |
| 602 | // Collect all function attrs for the node. |
| 603 | for (auto& iter : node.attrs()) { |
| 604 | if (iter.second.has_func()) { |
| 605 | VLOG(2) << "Found function attr for node " << node.name() << ": " |
| 606 | << iter.first << " = " << iter.second.func().name(); |
| 607 | results.emplace_back(AssociatedFunctionInfo::FunctionAttr( |
| 608 | iter.second.func().name(), iter.second.func().attr(), iter.first)); |
| 609 | } |
| 610 | } |
| 611 | } |
| 612 | return results; |
| 613 | } |
| 614 | |
| 615 | Status RewriteAssociatedFunction( |
| 616 | Graph* graph, Node* node, FunctionLibraryDefinition* fld, |
no test coverage detected