Function IndexSampleInner

paddle/phi/kernels/cpu/index_sample_kernel.cc:33–83

namespace phi {
template <typename T, typename Context, typename IndexT = int>
void IndexSampleInner(const Context &dev_ctx,
                      const DenseTensor &input,
                      const DenseTensor &index,
                      DenseTensor *output) {
  auto input_dims = input.dims();
  auto index_dims = index.dims();

  int64_t batch_size = input_dims[0];
  auto value_length = input_dims[1];
  auto index_length = index_dims[1];
  int64_t index_ids_num = index.numel();

  std::vector<T> input_vec;
  std::vector<IndexT> index_vec;
  TensorToVector(input, dev_ctx, &input_vec);
  TensorToVector<IndexT>(index, dev_ctx, &index_vec);

  std::vector<T> res(index_ids_num);
  for (int64_t i = 0; i < index_ids_num; i++) {
    int64_t b = floor(i / index_length);
    PADDLE_ENFORCE_GE(
        index_vec[i],
        0,
        errors::InvalidArgument(
            "Variable value (index) of OP(index_sample) "
            "expected >= 0 and < %ld, but got %ld. Please check input "
            "value.",
            value_length,
            index_vec[i]));
    PADDLE_ENFORCE_LT(
        index_vec[i],
        value_length,
        errors::InvalidArgument(
            "Variable value (index) of OP(index_sample) "
            "expected >= 0 and < %ld, but got %ld. Please check input "
            "value.",
            value_length,
            index_vec[i]));

    int64_t v_i = b * value_length + static_cast<int64_t>(index_vec[i]);
    T v = input_vec[v_i];
    VLOG(4) << "Index Sample: batch = " << b << " index = " << v_i
            << " value = " << v;
    res[i] = v;
  }

  auto ddim = make_ddim({batch_size, index_length});
  dev_ctx.template Alloc<T>(output);
  TensorFromVector(res, dev_ctx, output);
  output->Resize(ddim);
}

template <typename T, typename Context>
void IndexSampleKernel(const Context &dev_ctx,