}

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);