Verifies that the `wait_for` attribute on collective nodes matches `wait_for_map`.
| 62 | // Verifies that the `wait_for` attribute on collective nodes matches |
| 63 | // `wait_for_map`. |
| 64 | void VerifyAttrs( |
| 65 | const Graph& graph, |
| 66 | const std::unordered_map<string, std::vector<int32>> wait_for_map) { |
| 67 | for (const Node* node : graph.nodes()) { |
| 68 | if (node->IsCollective() || |
| 69 | wait_for_map.find(node->name()) == wait_for_map.end()) { |
| 70 | continue; |
| 71 | } |
| 72 | std::vector<int32> wait_for_actual; |
| 73 | TF_EXPECT_OK(GetNodeAttr(node->attrs(), "wait_for", &wait_for_actual)); |
| 74 | auto wait_for_expected = wait_for_map.at(node->name()); |
| 75 | EXPECT_THAT(wait_for_actual, UnorderedElementsAreArray(wait_for_expected)); |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | Node* CollectiveReduceNode(GraphDefBuilder* builder, Node* input, |
| 80 | const string& name, const string& device, |
no test coverage detected