Returns a list of tensors with the all-reduce sum across `tensors`. The computation is done with an all-reduce operation, so if only some of the returned tensors are evaluated then the computation will hang. Args: tensors: The input tensors across which to sum; must be assigned to
(tensors)
| 31 | |
| 32 | |
| 33 | def all_sum(tensors): |
| 34 | """Returns a list of tensors with the all-reduce sum across `tensors`. |
| 35 | |
| 36 | The computation is done with an all-reduce operation, so if only some of the |
| 37 | returned tensors are evaluated then the computation will hang. |
| 38 | |
| 39 | Args: |
| 40 | tensors: The input tensors across which to sum; must be assigned |
| 41 | to GPU devices. |
| 42 | |
| 43 | Returns: |
| 44 | List of tensors, each with the sum of the input tensors, where tensor i has |
| 45 | the same device as `tensors[i]`. |
| 46 | """ |
| 47 | return _apply_all_reduce('sum', tensors) |
| 48 | |
| 49 | |
| 50 | @ops.RegisterGradient('NcclAllReduce') |
nothing calls this directly
no test coverage detected