Returning numpy arrays via pybind11

I have a C ++ function calculating a large tensor that I would like to return Python as a NumPy array through pybind11 .

From the pybind11 documentation, it seems advisable to use STL unique_ptr . In the following example, the commented version works, while the given one compiles, but does not execute at runtime ("Cannot convert the return value of the function to Python type!").

Why is the smartpointer version not working? What is the canonical way to create and return a NumPy array?

PS: Due to the structure of the program and the size of the array, it is advisable not to copy the memory, but to create an array from the specified pointer. Memory ownership should be done by Python.

typedef typename py::array_t<double, py::array::c_style | py::array::forcecast> py_cdarray_t; // py_cd_array_t _test() std::unique_ptr<py_cdarray_t> _test() { double * memory = new double[3]; memory[0] = 11; memory[1] = 12; memory[2] = 13; py::buffer_info bufinfo ( memory, // pointer to memory buffer sizeof(double), // size of underlying scalar type py::format_descriptor<double>::format(), // python struct-style format descriptor 1, // number of dimensions { 3 }, // buffer dimensions { sizeof(double) } // strides (in bytes) for each index ); //return py_cdarray_t(bufinfo); return std::unique_ptr<py_cdarray_t>( new py_cdarray_t(bufinfo) ); } 
+7
c ++ python arrays numpy pybind11
source share
1 answer

A few comments (then a working implementation).

  • pybind11 C ++ wrappers around Python types (e.g. pybind11::object , pybind11::list and, in this case, pybind11::array_t<T> ) are really just wrappers around the base pointer to a Python object. In this respect, the role of wrapping a common pointer is already assumed, and therefore there is no point in wrapping that in unique_ptr : returning an object py::array_t<T> directly essentially essentially just returns a distinguished pointer.
  • pybind11::array_t can be constructed directly from the data pointer, so you can skip the intermediate step py::buffer_info and simply submit the form and go directly to the pybind11::array_t . The numpy array constructed in this way will not have its own data, it will just reference it (that is, the numpy owndata flag will be set to false).
  • Memory ownership may be tied to the life of a Python object, but you are still on the hook for the correct execution of the release. Pybind11 provides the py::capsule class to help you do just that. What you want to do is make the numpy array depend on this capsule as its parent class, specifying it as the base argument to array_t . This will make the numpy array reference it, saving it while the array itself is alive, and call the cleanup function when it no longer refers.
  • The c_style flag in earlier versions (before 2.2) only affected the new arrays, i.e. when no value pointer is passed. This was fixed in version 2.2 to also affect automatic steps if you specify only shapes, but not steps. It has no effect if you specify the steps directly (as in the example below).

So, combining the pieces, this code is a complete pybind11 module that demonstrates how you can accomplish what you are looking for (and includes some C ++ output that actually works correctly):

 #include <iostream> #include <pybind11/pybind11.h> #include <pybind11/numpy.h> namespace py = pybind11; PYBIND11_PLUGIN(numpywrap) { py::module m("numpywrap"); m.def("f", []() { // Allocate and initialize some data; make this big so // we can see the impact on the process memory use: constexpr size_t size = 100*1000*1000; double *foo = new double[size]; for (size_t i = 0; i < size; i++) { foo[i] = (double) i; } // Create a Python object that will free the allocated // memory when destroyed: py::capsule free_when_done(foo, [](void *f) { double *foo = reinterpret_cast<double *>(f); std::cerr << "Element [0] = " << foo[0] << "\n"; std::cerr << "freeing memory @ " << f << "\n"; delete[] foo; }); return py::array_t<double>( {100, 1000, 1000}, // shape {1000*1000*8, 1000*8, 8}, // C-style contiguous strides for double foo, // the data pointer free_when_done); // numpy array references this parent }); return m.ptr(); } 

Compiling and calling from Python shows that it works:

 >>> import numpywrap >>> z = numpywrap.f() >>> # the python process is now taking up a bit more than 800MB memory >>> z[1,1,1] 1001001.0 >>> z[0,0,100] 100.0 >>> z[99,999,999] 99999999.0 >>> z[0,0,0] = 3.141592 >>> del z Element [0] = 3.14159 freeing memory @ 0x7fd769f12010 >>> # python process memory size has dropped back down 
+15
source share

All Articles