mirror of
https://github.com/nomic-ai/kompute.git
synced 2026-05-12 01:19:58 +00:00
32 lines
945 B
C++
32 lines
945 B
C++
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
#include <kompute/Kompute.hpp>
|
|
|
|
namespace py = pybind11;
|
|
|
|
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;
|
|
#else
|
|
m.attr("__version__") = "dev";
|
|
#endif
|
|
}
|