The gradients for `all_sum`. Args: op: The `all_sum` `Operation` that we are differentiating. grad: Gradient with respect to the output of the `all_sum` op. Returns: The gradient with respect to the output of `all_sum`. Raises: LookupError: If `reduction` is not `sum`.
(op, grad)
| 49 | |
| 50 | @ops.RegisterGradient('NcclAllReduce') |
| 51 | def _all_sum_grad(op, grad): |
| 52 | """The gradients for `all_sum`. |
| 53 | |
| 54 | Args: |
| 55 | op: The `all_sum` `Operation` that we are differentiating. |
| 56 | grad: Gradient with respect to the output of the `all_sum` op. |
| 57 | |
| 58 | Returns: |
| 59 | The gradient with respect to the output of `all_sum`. |
| 60 | |
| 61 | Raises: |
| 62 | LookupError: If `reduction` is not `sum`. |
| 63 | """ |
| 64 | if op.get_attr('reduction') != b'sum': |
| 65 | raise LookupError('No gradient defined for NcclAllReduce except sum.') |
| 66 | |
| 67 | _check_device(grad, expected=op.device) |
| 68 | num_devices = op.get_attr('num_devices') |
| 69 | shared_name = op.get_attr('shared_name') + b'_grad' |
| 70 | |
| 71 | with ops.device(op.device): |
| 72 | return gen_nccl_ops.nccl_all_reduce( |
| 73 | input=grad, |
| 74 | reduction='sum', |
| 75 | num_devices=num_devices, |
| 76 | shared_name=shared_name) |
| 77 | |
| 78 | |
| 79 | def all_prod(tensors): |
nothing calls this directly
no test coverage detected