The gradients for input `Operation` of `reduce_sum`. Args: op: The `sum send` `Operation` that we are differentiating. grad: Gradient with respect to the output of the `reduce_sum` op. Returns: The gradient with respect to the input of `reduce_sum` op. Raises: LookupError: I
(op, grad)
| 148 | |
| 149 | @ops.RegisterGradient('NcclReduce') |
| 150 | def _reduce_sum_grad(op, grad): |
| 151 | """The gradients for input `Operation` of `reduce_sum`. |
| 152 | |
| 153 | Args: |
| 154 | op: The `sum send` `Operation` that we are differentiating. |
| 155 | grad: Gradient with respect to the output of the `reduce_sum` op. |
| 156 | |
| 157 | Returns: |
| 158 | The gradient with respect to the input of `reduce_sum` op. |
| 159 | |
| 160 | Raises: |
| 161 | LookupError: If the reduction attribute of op is not `sum`. |
| 162 | """ |
| 163 | if op.get_attr('reduction') != b'sum': |
| 164 | raise LookupError('No gradient defined for NcclReduce except sum.') |
| 165 | _check_device(grad, expected=op.device) |
| 166 | |
| 167 | with ops.device(op.device): |
| 168 | result = gen_nccl_ops.nccl_broadcast(input=grad, shape=grad.shape) |
| 169 | |
| 170 | return [result] * len(op.inputs) |
| 171 | |
| 172 | |
| 173 | def broadcast(tensor): |
nothing calls this directly
no test coverage detected