Returns the gradient of GatherV2 with batch dimensions.
(
params_shape, values, indices, batch_dims, gather_dim_size)
| 500 | |
| 501 | |
| 502 | def _BatchGatherGrad( |
| 503 | params_shape, values, indices, batch_dims, gather_dim_size): |
| 504 | """Returns the gradient of GatherV2 with batch dimensions.""" |
| 505 | |
| 506 | # Axis is the first non-batch dimension. |
| 507 | indices_size = array_ops.expand_dims(array_ops.size(indices), 0) |
| 508 | if batch_dims: |
| 509 | values_shape = array_ops.shape(values) |
| 510 | # Add the batch offsets to indices and flatten the batch dimensions. |
| 511 | outer_shape = values_shape[:batch_dims] |
| 512 | inner_shape = values_shape[batch_dims:][1:] |
| 513 | batch_size = gen_math_ops.prod(outer_shape, [0], False) |
| 514 | flat_values_shape = array_ops.concat([[-1], inner_shape], 0) |
| 515 | gather_dim_size *= batch_size |
| 516 | |
| 517 | indices = _GetBatchIndices(params_shape, indices, batch_dims) |
| 518 | |
| 519 | with warnings.catch_warnings(): |
| 520 | warnings.filterwarnings( |
| 521 | "ignore", |
| 522 | message="Converting sparse IndexedSlices to a dense Tensor.*") |
| 523 | values = array_ops.reshape(values, flat_values_shape) |
| 524 | |
| 525 | indices = array_ops.reshape(indices, indices_size) |
| 526 | params_grad = math_ops.unsorted_segment_sum(values, indices, gather_dim_size) |
| 527 | |
| 528 | if batch_dims: |
| 529 | # Put back the batch dimensions. |
| 530 | params_grad = array_ops.reshape( |
| 531 | params_grad, array_ops.concat([outer_shape, flat_values_shape], 0)) |
| 532 | |
| 533 | return params_grad |
| 534 | |
| 535 | |
| 536 | @ops.RegisterGradient("GatherV2") |
no test coverage detected