pybind11 与 numpy 的交互

1 #include <pybind11/pybind11.h> 2 #include <pybind11/numpy.h> 3 4 namespace py = pybind11; 5 6 py::array_t<double> add_arrays(py::array_t<double> input1, py::array_t<double> input2) { 7 py::buffer_info buf1 = input1.request(), buf2 = input2.request(); 8 9 if (buf1.ndim != 1 || buf2.ndim != 1) 10 throw std::runtime_error("Number of dimensions must be one"); 11 12 if (buf1.size != buf2.size) 13 throw std::runtime_error("Input shapes must match"); 14 15 /* No pointer is passed, so NumPy will allocate the buffer */ 16 auto result = py::array_t<double>(buf1.size); 17 18 py::buffer_info buf3 = result.request(); 19 20 double *ptr1 = (double *) buf1.ptr, 21 *ptr2 = (double *) buf2.ptr, 22 *ptr3 = (double *) buf3.ptr; 23 24 for (size_t idx = 0; idx < buf1.shape[0]; idx++) 25 ptr3[idx] = ptr1[idx] + ptr2[idx]; 26 27 return result; 28 } 29 30 PYBIND11_MODULE(test, m) { 31 m.def("add_arrays", &add_arrays, "Add two NumPy arrays"); 32 }