nomic-ai/kompute
1
0
Fork 0
mirror of https://github.com/nomic-ai/kompute.git synced 2026-05-11 17:09:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

added numpy() method

Browse Source
This commit is contained in:
---
2020-12-27 09:46:37 +01:00
parent e8b536c229
commit 46278eb0a9
1 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/src/main.cpp
View File

@@ -1,5 +1,6 @@
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <kompute/Kompute.hpp>
@@ -39,6 +40,20 @@ PYBIND11_MODULE(kp, m) {
return std::unique_ptr<kp::Tensor>(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<ssize_t> shape = { self.size() };
std::vector<ssize_t> strides = { sizeof(float) };
return py::array(py::buffer_info(
self.data().data(),
sizeof(float),
py::format_descriptor<float>::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) {