| 122 | } |
| 123 | |
| 124 | StatusOr<DeviceAssignment> ComputationPlacer::AssignDevices( |
| 125 | int replica_count, int computation_count) { |
| 126 | DeviceAssignment assignment(replica_count, computation_count); |
| 127 | for (int replica = 0; replica < replica_count; ++replica) { |
| 128 | for (int computation = 0; computation < computation_count; ++computation) { |
| 129 | TF_ASSIGN_OR_RETURN( |
| 130 | int device_id, |
| 131 | DeviceId(replica, computation, replica_count, computation_count)); |
| 132 | assignment(replica, computation) = device_id; |
| 133 | } |
| 134 | } |
| 135 | return std::move(assignment); |
| 136 | } |
| 137 | |
| 138 | /* static */ void ComputationPlacer::RegisterComputationPlacer( |
| 139 | se::Platform::Id platform_id, |
no outgoing calls