基于xtensor库实现numpy的take函数的技术实现与优化问询

功能需求 #

给定秩r≥1的data张量与秩q的indices张量，收集data指定axis维度（默认axis=0）中由indices索引的元素，拼接为秩为q+(r-1)的输出张量：

• 当axis=0时：设k = indices[i₀,...,i_{q-1}]，则output[i₀,...,i_{q-1},j₀,...,j_{r-2}] = input[k,j₀,...,j_{r-2}]
• 当axis=1时：设k = indices[i₀,...,i_{q-1}]，则output[j₀,i₀,...,i_{q-1},j₁,...,j_{r-2}] = input[j₀,k,j₁,...,j_{r-2}]

示例说明 #

示例1（axis=0） #

data = [
    [1.0, 1.2],
    [2.3, 3.4],
    [4.5, 5.7],
]
indices = [
    [0, 1],
    [1, 2],
]
output = [
    [
        [1.0, 1.2],
        [2.3, 3.4],
    ],
    [
        [2.3, 3.4],
        [4.5, 5.7],
    ],
]

示例2（axis=1） #

data = [
    [1.0, 1.2, 1.9],
    [2.3, 3.4, 3.9],
    [4.5, 5.7, 5.9],
]
indices = [
    [0, 2],
]
axis = 1,
output = [
        [[1.0, 1.9]],
        [[2.3, 3.9]],
        [[4.5, 5.9]],
]

当前实现与疑问 #

尝试过通过循环收集子数组视图到vector中但不知如何合并，也试过xt::tuple和xt::stack但编译失败。最终编写了如下模板函数，能运行示例，但不确定效率与正确性，寻求优化建议或更优实现方案：

template <class T>
xt::xarray<T> take(const xt::xarray<T>& arr, const xt::xarray<int>& indexes, int axis) {
  // TODO: make sure the implementation is correct
  std::vector<std::int64_t> new_shape{};
  new_shape.insert(new_shape.end(), arr.shape().begin(), arr.shape().begin() + axis);
  new_shape.insert(new_shape.end(), indexes.shape().begin(), indexes.shape().end());
  new_shape.insert(new_shape.end(), arr.shape().begin() + axis + 1, arr.shape().end());
  xt::xarray<T> new_values = xt::zeros<T>(new_shape);

  int counter = 0;
  for (const int index : indexes) {
    xt::xstrided_slice_vector data_slices{};
    xt::xstrided_slice_vector new_arr_slices{};
    for (int i = 0; i < axis; ++i) {
      data_slices.emplace_back(xt::all());
      new_arr_slices.emplace_back(xt::all());
    }
    for (size_t i = 1; i < indexes.dimension(); ++i) {
      new_arr_slices.emplace_back(0);
    }
    data_slices.emplace_back(index);
    new_arr_slices.emplace_back(counter++);
    data_slices.emplace_back(xt::ellipsis());
    new_arr_slices.emplace_back(xt::ellipsis());

    auto data_slice = xt::strided_view(arr, data_slices);
    auto new_slice = xt::strided_view(new_values, new_arr_slices);
    for (size_t j = 0; j < data_slice.size(); ++j) {
      new_slice(j) = data_slice(j);
    }
  }

  return new_values;
}

内容的提问来源于stack exchange，提问作者Tend