| 1045 | } |
| 1046 | |
| 1047 | Node* InternalFindUniqueNode(Node* n) { |
| 1048 | std::unordered_set<Node*> pushed; |
| 1049 | std::queue<Node*> q; |
| 1050 | q.push(n); |
| 1051 | pushed.insert(n); |
| 1052 | while (!q.empty()) { |
| 1053 | Node* curr = q.front(); |
| 1054 | q.pop(); |
| 1055 | if (curr->op_def().name() == "Unique" || |
| 1056 | curr->op_def().name() == "UniqueV2") { |
| 1057 | return curr; |
| 1058 | } |
| 1059 | |
| 1060 | for (const Edge* edge : curr->in_edges()) { |
| 1061 | if (edge->IsControlEdge()) continue; |
| 1062 | if (pushed.find(edge->src()) == pushed.end()) { |
| 1063 | pushed.insert(edge->src()); |
| 1064 | q.push(edge->src()); |
| 1065 | } |
| 1066 | } |
| 1067 | } |
| 1068 | |
| 1069 | return nullptr; |
| 1070 | } |
| 1071 | |
| 1072 | Node* CheckUniqueNode(std::vector<Node*> gather_nodes) { |
| 1073 | Node* unique_node = nullptr; |