// 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.
Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
                          const Output& reduction_axes);

}  // namespace tensorflow

using tensorflow::ops::Const;
using tensorflow::ops::DynamicStitch;
using tensorflow::ops::Mod;
using tensorflow::ops::OnesLike;
using tensorflow::ops::Range;
using tensorflow::ops::Size;

Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
                          const Output& reduction_axes) {
  auto zero = Const(scope, 0);
  auto one = Const(scope, 1);

  // Running example in comments

  // [[g1, g1, g1],
  //  [g2, g2, g2]]

  // input_shape = [2, 3]
  auto input_shape = Shape(scope, op.input(0));

  // output_shape_kept_dims = [2, 1]
  auto output_shape_kept_dims =
      ReducedShapeHelper(scope, input_shape, op.input(1));

  // This step "flips" any 1s with values from the input_shape, and
  // replaces remaining entries with 1. This creates a shape that

    // Obtain the shape of the output, as if keepdims=True on reduce sum. E.g.
    // for the equation "abcd->ac" with input shape [2,5,3,4], we get the
    // reduced shape [2,1,3,1].
    auto reduced_shape = ReducedShapeHelper(scope, input_shape, reduced_axes);
    // Reshaping the gradient (wrt "ac") to [2,1,3,1] and broadcasting it to
    // the shape [2,5,3,4] results in the gradient wrt "abcd".