Gradient for RaggedGather op.
(op, *grads)
| 270 | #=============================================================================== |
| 271 | @ops.RegisterGradient('RaggedGather') |
| 272 | def _ragged_gather_grad(op, *grads): |
| 273 | """Gradient for RaggedGather op.""" |
| 274 | param_nested_splits = op.inputs[:-2] |
| 275 | param_inner_values = op.inputs[-2] |
| 276 | indices = op.inputs[-1] |
| 277 | grad_inner_values = grads[-1] |
| 278 | |
| 279 | # For each row in `params`, find the range of values in `params.inner_values` |
| 280 | # that is covered by that row. In particular, the values in row `i` are |
| 281 | # `param_inner_values[combined_splits[i]:combined_splits[i+1]`. |
| 282 | combined_splits = param_nested_splits[0] |
| 283 | for row_splits in param_nested_splits[1:]: |
| 284 | combined_splits = array_ops.gather(row_splits, combined_splits) |
| 285 | |
| 286 | # The outer dimensions of `indices` correspond 1:1 with the outer dimensions |
| 287 | # of `ragged_grad` that are encoded by `grad_nested_splits`. Thus, the |
| 288 | # flattened `indices` correspond 1:1 with `grad_inner_values`. |
| 289 | flat_indices = array_ops.reshape(indices, [-1]) |
| 290 | |
| 291 | # Build an IndexedSlices where the values are taken from `flat_grad`. |
| 292 | grad_indices = ragged_math_ops.range( |
| 293 | array_ops.gather(combined_splits, flat_indices), |
| 294 | array_ops.gather(combined_splits[1:], flat_indices)).values |
| 295 | |
| 296 | param_inner_values_grad = indexed_slices.IndexedSlices( |
| 297 | values=grad_inner_values, indices=grad_indices, |
| 298 | dense_shape=array_ops.shape(param_inner_values)) |
| 299 | return [None for _ in param_nested_splits] + [param_inner_values_grad, None] |