Call nccl allreduce.
()
| 215 | shared_name = _get_shared_name() |
| 216 | |
| 217 | def _all_reduce(): |
| 218 | """Call nccl allreduce.""" |
| 219 | res = [] |
| 220 | for t in tensors: |
| 221 | _check_device(t) |
| 222 | with ops.device(t.device): |
| 223 | res.append( |
| 224 | gen_nccl_ops.nccl_all_reduce( |
| 225 | input=t, |
| 226 | reduction=reduction, |
| 227 | num_devices=len(tensors), |
| 228 | shared_name=shared_name)) |
| 229 | return res |
| 230 | |
| 231 | if context.executing_eagerly(): |
| 232 | # Nccl ops will block unless they are executed concurrently such as in a |
no test coverage detected