TODO(b/77601805): add tests for associated function related stuff.
| 557 | |
| 558 | // TODO(b/77601805): add tests for associated function related stuff. |
| 559 | bool HasAssociatedFunction(const NodeDef& node_def, |
| 560 | const FunctionLibraryDefinition* fld) { |
| 561 | if (fld->Contains(node_def.op())) { |
| 562 | return true; |
| 563 | } |
| 564 | |
| 565 | if (node_def.op() == FunctionLibraryDefinition::kGradientOp) { |
| 566 | // Gradient op has "f" attr, which is set to the function we are getting |
| 567 | // gradient for. We need to functionalize the gradient function. |
| 568 | return true; |
| 569 | } |
| 570 | |
| 571 | if (node_def.op() == "XlaHostCompute") { |
| 572 | // XlaHostCompute has "shape_inference_graph" func attr, but that's not |
| 573 | // related to graph execution. |
| 574 | return false; |
| 575 | } |
| 576 | |
| 577 | for (const auto& iter : node_def.attr()) { |
| 578 | if (iter.second.has_func()) { |
| 579 | return true; |
| 580 | } |
| 581 | } |
| 582 | |
| 583 | return false; |
| 584 | } |
| 585 | |
| 586 | std::vector<AssociatedFunctionInfo> GetAssociatedFunctions( |
| 587 | const Node& node, const FunctionLibraryDefinition* fld) { |