Common code for SparseDenseCwise{Mul,Div} gradients.
(op, grad, is_mul)
| 208 | |
| 209 | |
| 210 | def _SparseDenseCwiseMulOrDivGrad(op, grad, is_mul): |
| 211 | """Common code for SparseDenseCwise{Mul,Div} gradients.""" |
| 212 | x_indices = op.inputs[0] |
| 213 | x_shape = op.inputs[2] |
| 214 | y = op.inputs[3] |
| 215 | |
| 216 | y_shape = math_ops.cast(array_ops.shape(y), dtypes.int64) |
| 217 | num_added_dims = array_ops.expand_dims( |
| 218 | array_ops.size(x_shape) - array_ops.size(y_shape), 0) |
| 219 | augmented_y_shape = array_ops.concat( |
| 220 | [array_ops.ones(num_added_dims, ops.dtypes.int64), y_shape], 0) |
| 221 | |
| 222 | scaling = x_shape // augmented_y_shape |
| 223 | scaled_indices = x_indices // scaling |
| 224 | scaled_indices = array_ops.slice(scaled_indices, |
| 225 | array_ops.concat([[0], num_added_dims], 0), |
| 226 | [-1, -1]) |
| 227 | dense_vals = array_ops.gather_nd(y, scaled_indices) |
| 228 | |
| 229 | if is_mul: |
| 230 | dx = grad * dense_vals |
| 231 | dy_val = grad * op.inputs[1] |
| 232 | else: |
| 233 | dx = grad / dense_vals |
| 234 | dy_val = grad * (-op.inputs[1] / math_ops.square(dense_vals)) |
| 235 | # indices can repeat after scaling, so we can't use sparse_to_dense(). |
| 236 | dy = sparse_ops.sparse_add( |
| 237 | array_ops.zeros_like(y), |
| 238 | sparse_tensor.SparseTensor(scaled_indices, dy_val, y_shape)) |
| 239 | |
| 240 | # (sp_indices, sp_vals, sp_shape, dense) |
| 241 | return (None, dx, None, dy) |
| 242 | |
| 243 | |
| 244 | @ops.RegisterGradient("SparseDenseCwiseMul") |