nomic-ai/kompute
1
0
Fork 0
mirror of https://github.com/nomic-ai/kompute.git synced 2026-05-24 06:54:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added base capabilities for tensor in python

Browse Source
This commit is contained in:
Alejandro Saucedo
2020-11-01 15:14:22 +00:00
parent 9559c79eee
commit 0e9ba00b71
4 changed files with 35 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/src/main.cpp
View File

@@ -1,9 +1,27 @@
#include <pybind11/pybind11.h>
#include "kompute/Kompute.hpp"
#include <pybind11/stl.h>
#include <kompute/Kompute.hpp>
namespace py = pybind11;
PYBIND11_MODULE(cmake_example, m) {
PYBIND11_MODULE(komputepy, m) {
py::enum_<kp::Tensor::TensorTypes>(m, "TensorTypes")
.value("eDevice", kp::Tensor::TensorTypes::eDevice)
.value("eStaging", kp::Tensor::TensorTypes::eStaging)
.value("eStorage", kp::Tensor::TensorTypes::eStorage)
.export_values();
py::class_<kp::Tensor>(m, "Tensor")
.def(py::init(
[](const std::vector<float>& data) {
return std::unique_ptr<kp::Tensor>(new kp::Tensor(data));
}))
.def(py::init(
[](const std::vector<float>& data, kp::Tensor::TensorTypes tensorTypes) {
return std::unique_ptr<kp::Tensor>(new kp::Tensor(data, tensorTypes));
}));
#ifdef VERSION_INFO
m.attr("__version__") = VERSION_INFO;