Function BatchGatherGrad

tensorflow/cc/gradients/array_grad.cc:649–676

}

Output BatchGatherGrad(const Scope& scope, Output params_shape, Output values,
                       Output indices, int batch_dims, Output gather_dim_size) {
  // Axis is the first non-batch dimension.
  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
  Output outer_shape, flat_values_shape;
  if (batch_dims != 0) {
    auto values_shape = Shape(scope, values);
    // Add the batch offsets to indices and flatten the batch dimensions.
    outer_shape = Slice(scope, values_shape, {0}, {batch_dims});
    auto inner_shape =
        Slice(scope, Slice(scope, values_shape, {batch_dims}, {-1}), {1}, {-1});
    auto batch_size = Prod(scope, outer_shape, /*axis=*/0);
    flat_values_shape = Concat(scope, {{-1}, inner_shape}, /*axis=*/0);
    gather_dim_size = Multiply(scope, gather_dim_size, batch_size);
    indices = GetBatchIndices(scope, params_shape, indices, batch_dims);
    values = Reshape(scope, values, flat_values_shape);
  }

  indices = Reshape(scope, indices, indices_size);
  Output params_grad =
      UnsortedSegmentSum(scope, values, indices, gather_dim_size);

  if (batch_dims != 0) {
    // Put back the batch dimensions.
    params_grad = Reshape(scope, params_grad, params_shape);
  }
  return params_grad;
}

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