diff --git a/python/src/main.cpp b/python/src/main.cpp index 6f3431661..871354a3c 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -1,5 +1,6 @@ #include #include +#include #include @@ -39,6 +40,20 @@ PYBIND11_MODULE(kp, m) { return std::unique_ptr(new kp::Tensor(data, tensorTypes)); }), "Initialiser with list of data components and tensor GPU memory type.") .def("data", &kp::Tensor::data, DOC(kp, Tensor, data)) + .def("numpy", [](kp::Tensor& self){ + ssize_t ndim = 1; + std::vector shape = { self.size() }; + std::vector strides = { sizeof(float) }; + + return py::array(py::buffer_info( + self.data().data(), + sizeof(float), + py::format_descriptor::format(), + ndim, + shape, + strides + )); + }, "Returns stored data as a new numpy array.") .def("__getitem__", [](kp::Tensor &self, size_t index) -> float { return self.data()[index]; }, "When only an index is necessary") .def("__setitem__", [](kp::Tensor &self, size_t index, float value) {