Helper function for reduction ops. input_shape: 1-D Tensor, the shape of the Tensor being reduced. axes: 1-D Tensor, the reduction axes. Note that the reduction indices are in the range -rank(input_shape), rank(input_shape) returns a 1-D Tensor, the output shape as if keep_dims were set to True.
| 642 | // -rank(input_shape), rank(input_shape) |
| 643 | // returns a 1-D Tensor, the output shape as if keep_dims were set to True. |
| 644 | Output ReducedShapeHelper(const Scope& scope, const Output& input_shape, |
| 645 | const Output& reduction_axes) { |
| 646 | auto zero = Const(scope, 0); |
| 647 | auto one = Const(scope, 1); |
| 648 | |
| 649 | // Running example in comments |
| 650 | // input_shape = [2, 3, 5, 7] |
| 651 | // axes = [1, 2] |
| 652 | // The result (a shape after a reduction with keep_dims=True) |
| 653 | // [2, 1, 1, 7] |
| 654 | // |
| 655 | // We can treat each entry in axes as an index into input_shape that |
| 656 | // should be replaced by 1. |
| 657 | // We use DynamicStitch to do this. |
| 658 | |
| 659 | // input_rank = 4 |
| 660 | auto input_rank = Size(scope, input_shape); |
| 661 | |
| 662 | // Normalize any negative indices in the reduction_axes to positive |
| 663 | // values. |
| 664 | auto axes = Mod(scope, Add(scope, reduction_axes, input_rank), input_rank); |
| 665 | |
| 666 | // This [0..input_rank) range of integers is used in DynamicStitch to |
| 667 | // first copy input_shape to the result. |
| 668 | // input_rank_range = [0, 1, 2, 3] |
| 669 | auto input_rank_range = Range(scope, zero, input_rank, one); |
| 670 | |
| 671 | // A 1-filled tensor with the same shape as axes. DynamicStitch will |
| 672 | // merge these 1s (using axes for indices) to the correct |
| 673 | // position in the result. |
| 674 | // axes_ones = [1, 1] |
| 675 | auto axes_ones = OnesLike(scope, axes); |
| 676 | |
| 677 | // using DynamicStitch: |
| 678 | // indices = { input_rank_range, axes } |
| 679 | // = { [0, 1, 2, 3], [1, 2] } |
| 680 | // data = { input_shape, axes_ones } |
| 681 | // = { [2, 3, 5, 7], [1, 1] } |
| 682 | // The input_rank_range entry in indices first replicates the |
| 683 | // input_shape to the result. |
| 684 | // The axes entry in indices then moves a 1 to each of its entries, |
| 685 | // resulting in |
| 686 | // [2, 1, 1, 7] |
| 687 | std::vector<Output> indices = {input_rank_range, axes}; |
| 688 | std::vector<Output> data = {input_shape, axes_ones}; |
| 689 | return DynamicStitch(scope, indices, data); |
| 690 | } |
| 691 | |
| 692 | // SumGradHelper returns the gradient for the Sum operator, and is used |
| 693 | // by SumGrad and MeanGrad. |