Function ConcatGradHelper

tensorflow/cc/gradients/array_grad.cc:508–549

REGISTER_GRADIENT_OP("Slice", SliceGrad);

Status ConcatGradHelper(const Scope& scope, const Operation& op,
                        const std::vector<Output>& grad_inputs,
                        std::vector<Output>* grad_outputs,
                        int start_value_index, int end_value_index,
                        int dim_index) {
  if (end_value_index >= op.num_inputs()) {
    return errors::Internal("Invalid input index");
  }
  std::vector<Output> inputs;
  for (int i = start_value_index; i < end_value_index; ++i) {
    inputs.push_back(op.input(i));
  }

  auto shapes = ShapeN(scope, inputs);
  const auto unique_name = scope.GetUniqueNameForOp("ConcatOffset");
  auto builder =
      ::tensorflow::NodeBuilder(unique_name, "ConcatOffset")
          .Input(::tensorflow::ops::AsNodeOut(scope, op.input(dim_index)))
          .Input(::tensorflow::ops::AsNodeOutList(scope, shapes.output));
  scope.UpdateBuilder(&builder);
  ::tensorflow::Node* concat_offset_node;
  scope.UpdateStatus(builder.Finalize(scope.graph(), &concat_offset_node));
  scope.UpdateStatus(scope.DoShapeInference(concat_offset_node));
  if (concat_offset_node->num_outputs() != inputs.size()) {
    return errors::Internal("ConcatOffset has invalid output count");
  }
  if (grad_inputs.size() != 1) {
    return errors::InvalidArgument("Concat grad should have 1 input");
  }

  // For each dx[i], we take a slice of dy. The offset and size of the
  // slice is given by offset[i] and shape[i].
  const Output& dy = grad_inputs[0];
  for (int i = 0; i < inputs.size(); ++i) {
    grad_outputs->push_back(
        Slice(scope, dy, Output(concat_offset_node, i), shapes.output[i]));
  }

  // Insert a NoGradient for the axis.
  grad_outputs->insert(grad_outputs->begin() + dim_index, NoGradient());
  return scope.status();
}

Status ConcatV2Grad(const Scope& scope, const Operation& op,
                    const std::vector<Output>& grad_inputs,