Returns a tensor that can be efficiently transferred to other devices. Args: tensor: The tensor to send; must be assigned to a GPU device. Returns: A tensor with the value of `src_tensor`, which can be used as input to ops on other GPU devices.
(tensor)
| 171 | |
| 172 | |
| 173 | def broadcast(tensor): |
| 174 | """Returns a tensor that can be efficiently transferred to other devices. |
| 175 | |
| 176 | Args: |
| 177 | tensor: The tensor to send; must be assigned to a GPU device. |
| 178 | |
| 179 | Returns: |
| 180 | A tensor with the value of `src_tensor`, which can be used as input to |
| 181 | ops on other GPU devices. |
| 182 | """ |
| 183 | _check_device(tensor) |
| 184 | |
| 185 | with ops.device(tensor.device): |
| 186 | return gen_nccl_ops.nccl_broadcast(input=tensor, shape=tensor.shape) |
| 187 | |
| 188 | |
| 189 | @ops.RegisterGradient('NcclBroadcast') |
nothing calls this directly
no test coverage detected