Function GatherV2Grad

tensorflow/cc/gradients/array_grad.cc:678–760

}

Status GatherV2Grad(const Scope& scope, const Operation& op,
                    const std::vector<Output>& grad_inputs,
                    std::vector<Output>* grad_outputs) {
  if (op.num_inputs() != 3) {
    return errors::InvalidArgument("Gather requires 3 inputs");
  }
  if (grad_inputs.size() != 1) {
    return errors::InvalidArgument("Gather grad requires 1 grad input");
  }

  // params can be large, so colocate the shape calculation with it.
  // params can be very large for sparse model, array_ops.shape raises
  // exception on the Windows platform when any dimension is larger than
  // int32. params_shape is not used in optimizer apply_sparse gradients,
  // so it's fine to convert it back to int32 regardless of truncation.
  auto params = op.input(0);
  auto colocate_scope = scope.ColocateWith(params);
  Shape::Attrs shape_attrs;
  shape_attrs.out_type_ = DT_INT64;
  auto params_shape64 = Shape(colocate_scope, params, shape_attrs);
  Output params_shape = Cast(colocate_scope, params_shape64, DT_INT32);

  auto indices = op.input(1);
  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
  auto axis = op.input(2);
  auto axis_expand = ExpandDims(scope, axis, 0);

  int batch_dims;
  TF_RETURN_IF_ERROR(
      GetNodeAttr(op.node()->attrs(), "batch_dims", &batch_dims));
  if (batch_dims < 0) {
    // TODO(bdodson): Figure out if we can find the param rank here, like the
    // python implementation does.
    return errors::InvalidArgument(
        "C++ GatherV2 gradient does not support negative batch_dims.");
  }

  // Handle axis by transposing the axis dimension to be the first non-batch
  // dimension, compute the gradient and transpose the result back.
  auto outer_shape = Slice(scope, params_shape, {0}, axis_expand);
  auto inner_shape =
      Slice(scope, Slice(scope, params_shape, axis_expand, {-1}), {1}, {-1});
  auto values_shape = Concat(scope, {outer_shape, {-1}, inner_shape}, 0);
  auto values_dims = Size(scope, values_shape);
  auto axis_dims = Size(scope, outer_shape);

  Output outer_batches_indices = Range(scope, 0, batch_dims, /*delta=*/1);
  Output batch_axis_indices = Range(scope, batch_dims, axis_dims, /*delta=*/1);
  Output inner_axes_indices =
      Range(scope, Add(scope, axis_dims, 1), values_dims, /*delta=*/1);
  Output axis_dims_expand = ExpandDims(scope, axis_dims, 0);

  auto values = Reshape(scope, grad_inputs[0], values_shape);

  // Move values[axis] up to values[batch_dims]
  Output transpose_dims = Concat(scope,
                                 {outer_batches_indices, axis_dims_expand,
                                  batch_axis_indices, inner_axes_indices},