如何用pybind11绑定一个以numpy.array（）作为参数（例如，形状为（10，10，3））的函数？

Question

I would like to write a function that can take a multidimensional numpy array, not just 2D.

void compute(Eigen::Ref<Eigen::MatrixXd> array3d) {
    // change the array in-place
    // ...
}

or

Eigen::MatrixXd &compute() {
    // create array
    // ...
    // and return it
}

I am using Eigen here just to depict the goal, I believe Eigen does not support 3D, or more dimensions, arrays.

I appreciate your feedback and patience since I am not familiar with Pybind11 nor Eigen.

Answer 1

From the pybind information, you can extract the dimension information.

For instance, this is what I do inside Audio ToolKit with m the current Python module you want to build:

py::class_<MyClass>(m,"name")
.def("set_pointer", [](MyClass& instance, const py::array_t<DataType>& array)
{
  gsl::index channels = 1;
  gsl::index size = array.shape(0);
  if(array.ndim() == 2)
  {
    channels = array.shape(0);
    size = array.shape(1);
  }
  // Call using array.data() and possibly add more dimension information, this is specific to my use case
  instance.set_pointer(array.data(), channels, size);
});

From this, you can create the Eigen::Map call instead to create an Eigen-like matrix that you can use in your templated code.

Basically, pybind11 allows you to create a lambda where you can create your wrapper for your use case. The same works for return, you can get the Eigen class, create a pybind array that you populate with the Eigen data.

Eigen has the Tensor class that you can use as well.