The backward operator for the SparseAdd op. The SparseAdd op calculates A + B, where A, B, and the sum are all represented as `SparseTensor` objects. This op takes in the upstream gradient w.r.t. non-empty values of the sum, and outputs the gradients w.r.t. the non-empty values of A and B.
(op, *grads)
| 64 | |
| 65 | @ops.RegisterGradient("SparseAdd") |
| 66 | def _SparseAddGrad(op, *grads): |
| 67 | """The backward operator for the SparseAdd op. |
| 68 | |
| 69 | The SparseAdd op calculates A + B, where A, B, and the sum are all represented |
| 70 | as `SparseTensor` objects. This op takes in the upstream gradient w.r.t. |
| 71 | non-empty values of the sum, and outputs the gradients w.r.t. the non-empty |
| 72 | values of A and B. |
| 73 | |
| 74 | Args: |
| 75 | op: the SparseAdd op |
| 76 | *grads: the incoming gradients, one element per output of `op` |
| 77 | |
| 78 | Returns: |
| 79 | Gradient for each of the 6 input tensors of SparseAdd: |
| 80 | (a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh) |
| 81 | The gradients for the indices, shapes, and the threshold are None. |
| 82 | """ |
| 83 | val_grad = grads[1] |
| 84 | a_indices = op.inputs[0] |
| 85 | b_indices = op.inputs[3] |
| 86 | sum_indices = op.outputs[0] |
| 87 | # NOTE: we do not need to take `thresh` into account, since it simply affects |
| 88 | # the non-zero elements of the sum, and we will peek into `sum_indices` in the |
| 89 | # gradient op. |
| 90 | |
| 91 | a_val_grad, b_val_grad = gen_sparse_ops.sparse_add_grad( |
| 92 | val_grad, a_indices, b_indices, sum_indices) |
| 93 | a_val_grad.set_shape(op.inputs[1].get_shape()) |
| 94 | b_val_grad.set_shape(op.inputs[4].get_shape()) |
| 95 | # (a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh) |
| 96 | return (None, a_val_grad, None, None, b_val_grad, None, None) |
| 97 | |
| 98 | |
| 99 | @ops.RegisterGradient("SparseTensorDenseAdd") |