The gradients for input `Operation` of `broadcast`. Args: op: The `broadcast send` `Operation` that we are differentiating. accumulated_grad: Accumulated gradients with respect to the output of the `broadcast` op. Returns: Gradients with respect to the input of `broadcast`.
(op, accumulated_grad)
| 188 | |
| 189 | @ops.RegisterGradient('NcclBroadcast') |
| 190 | def _broadcast_grad(op, accumulated_grad): |
| 191 | """The gradients for input `Operation` of `broadcast`. |
| 192 | |
| 193 | Args: |
| 194 | op: The `broadcast send` `Operation` that we are differentiating. |
| 195 | accumulated_grad: Accumulated gradients with respect to the output of the |
| 196 | `broadcast` op. |
| 197 | |
| 198 | Returns: |
| 199 | Gradients with respect to the input of `broadcast`. |
| 200 | """ |
| 201 | # Grab inputs of accumulated_grad and replace accumulation with reduce_sum. |
| 202 | grads = [t for t in accumulated_grad.op.inputs] |
| 203 | for t in grads: |
| 204 | _check_device(t) |
| 205 | |
| 206 | with ops.device(op.device): |
| 207 | return gen_nccl_ops.nccl_reduce(input=grads, reduction='sum') |
| 208 | |
| 209 | |
| 210 | def _apply_all_reduce(reduction, tensors): |
nothing calls this directly
no test coverage detected