Function GetBatchIndices

tensorflow/cc/gradients/array_grad.cc:621–647

Adds the batch offsets to the given indices and returns the results.

// Adds the batch offsets to the given indices and returns the results.
Output GetBatchIndices(const Scope& scope, const Output& params_shape,
                       const Output& indices, int batch_dims) {
  Output batch_indices = indices;
  auto indices_ndims = Rank(scope, indices);
  auto casted_params_shape = Cast(scope, params_shape, indices.type());
  Output accum_dim_value = ConstHelper(scope, 1, indices.type());
  for (int dim = batch_dims; dim > 0; dim--) {
    Output dim_value = Slice(scope, casted_params_shape, {dim - 1}, {1});
    accum_dim_value = Multiply(scope, accum_dim_value,
                               Slice(scope, casted_params_shape, {dim}, {1}));
    auto start = ConstHelper(scope, 0, indices.type());
    auto step = ConstHelper(scope, 1, indices.type());
    Output dim_indices = Range(scope, start, Squeeze(scope, dim_value), step);
    dim_indices = Multiply(scope, dim_indices, accum_dim_value);
    auto one = Cast(scope, Const(scope, {1}), indices.type());
    auto dim_shape = Concat(
        scope,
        {Output(Tile(scope, one, Const(scope, {dim - 1}))), dim_value,
         Output(Tile(scope, one,
                     ExpandDims(scope, Sub(scope, indices_ndims, dim), 0)))},
        /*axis=*/0);
    batch_indices =
        Add(scope, batch_indices, Reshape(scope, dim_indices, dim_shape));
  }

  return batch_indices;
}

Output BatchGatherGrad(const Scope& scope, Output params_shape, Output values,
                       Output indices, int batch_dims, Output gather_dim_size) {