Makes a DeviceAssignment assigning replica-id i to devices[i].
| 150 | |
| 151 | // Makes a DeviceAssignment assigning replica-id i to devices[i]. |
| 152 | DeviceAssignment MakeDeviceAssn(std::vector<int64> devices) { |
| 153 | DeviceAssignment assn(/*replica_count=*/devices.size(), |
| 154 | /*computation_count=*/1); |
| 155 | for (int64 i = 0; i < devices.size(); ++i) { |
| 156 | assn(i, 0) = devices[i]; |
| 157 | } |
| 158 | return assn; |
| 159 | } |
| 160 | |
| 161 | // Shorter alias for this function. |
| 162 | absl::flat_hash_set<GlobalDeviceId> OpenNcclChannels() { |