Gradient for GatherV2 op.
(op, grad)
| 535 | |
| 536 | @ops.RegisterGradient("GatherV2") |
| 537 | def _GatherV2Grad(op, grad): |
| 538 | """Gradient for GatherV2 op.""" |
| 539 | # params can be large, so colocate the shape calculation with it. |
| 540 | # |
| 541 | # params can be very large for sparse model, array_ops.shape raises |
| 542 | # exception on the Windows platform when any dimension is larger than |
| 543 | # int32. params_shape is not used in optimizer apply_sparse gradients, |
| 544 | # so it's fine to convert it back to int32 regardless of truncation. |
| 545 | params = op.inputs[0] |
| 546 | with ops.colocate_with(params): |
| 547 | params_shape = array_ops.shape(params, out_type=ops.dtypes.int64) |
| 548 | params_shape = math_ops.cast(params_shape, dtypes.int32) |
| 549 | |
| 550 | indices = op.inputs[1] |
| 551 | indices_size = array_ops.expand_dims(array_ops.size(indices), 0) |
| 552 | axis = op.inputs[2] |
| 553 | axis_static = tensor_util.constant_value(axis) |
| 554 | batch_dims = int(op.get_attr("batch_dims")) |
| 555 | |
| 556 | if batch_dims < 0: |
| 557 | batch_dims += indices.shape.ndims |
| 558 | |
| 559 | # For axis 0 gathers, build an appropriately shaped IndexedSlices. |
| 560 | if axis_static == 0: |
| 561 | if context.executing_eagerly(): |
| 562 | params_tail_shape = params_shape.cpu()[1:] |
| 563 | else: |
| 564 | params_tail_shape = params_shape[1:] |
| 565 | values_shape = array_ops.concat([indices_size, params_tail_shape], 0) |
| 566 | with warnings.catch_warnings(): |
| 567 | warnings.filterwarnings( |
| 568 | "ignore", |
| 569 | message="Converting sparse IndexedSlices to a dense Tensor.*") |
| 570 | values = array_ops.reshape(grad, values_shape) |
| 571 | indices = array_ops.reshape(indices, indices_size) |
| 572 | params_grad = ops.IndexedSlices(values, indices, params_shape) |
| 573 | else: |
| 574 | # Handle axis by transposing the axis dimension to be the first non-batch |
| 575 | # dimension, compute the gradiend and transpose the result back. |
| 576 | outer_shape = params_shape[:axis] |
| 577 | inner_shape = params_shape[axis:][1:] |
| 578 | values_shape = array_ops.concat([outer_shape, [-1], inner_shape], 0) |
| 579 | |
| 580 | values_dims = array_ops.size(values_shape) |
| 581 | axis_dims = array_ops.size(outer_shape) |
| 582 | |
| 583 | outer_batches_indices = math_ops.range(batch_dims) |
| 584 | batch_axis_indices = math_ops.range(batch_dims, axis_dims) |
| 585 | inner_axes_indices = math_ops.range(axis_dims + 1, values_dims) |
| 586 | |
| 587 | with warnings.catch_warnings(): |
| 588 | warnings.filterwarnings( |
| 589 | "ignore", |
| 590 | message="Converting sparse IndexedSlices to a dense Tensor.*") |
| 591 | values = array_ops.reshape(grad, values_shape) |
| 592 | |
| 593 | # Move values[axis] up to values[batch_dims] |
| 594 | transpose_dims = array_ops.concat( |
nothing calls this directly
no test coverage detected