| 335 | } |
| 336 | |
| 337 | Status CallGraph::VisitNodes(const VisitorFunction& visitor_func, |
| 338 | bool visit_unreachable_nodes) const { |
| 339 | absl::flat_hash_set<const CallGraphNode*> visited; |
| 340 | if (visit_unreachable_nodes) { |
| 341 | // Traverse from all roots in the call graph. |
| 342 | for (const CallGraphNode& node : nodes()) { |
| 343 | if (node.callers().empty()) { |
| 344 | TF_RETURN_IF_ERROR(VisitNodesInternal(visitor_func, node, &visited)); |
| 345 | } |
| 346 | } |
| 347 | } else { |
| 348 | // Traverse only from the entry computation. |
| 349 | TF_RETURN_IF_ERROR(VisitNodesInternal( |
| 350 | visitor_func, GetNode(module_->entry_computation()), &visited)); |
| 351 | } |
| 352 | |
| 353 | return Status::OK(); |
| 354 | } |
| 355 | |
| 356 | bool CallGraph::IsFlattened() const { |
| 357 | for (const CallGraphNode& node : nodes_) { |