L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Switch to nanobind (#839)

Browse Source 
* mostly builds

* most tests pass

* fix circle build

* add back buffer protocol

* includes

* fix for py38

* limit to cpu device

* include

* fix stubs

* move signatures for docs

* stubgen + docs fix

* doc for compiled function, comments
This commit is contained in:
Awni Hannun
2024-03-18 20:12:25 -07:00
committed by GitHub
parent d39ed54f8e
commit 9a8ee00246
34 changed files with 2343 additions and 2344 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/device.cpp
+24 -18
View File
@@ -1,32 +1,34 @@
// Copyright © 2023 Apple Inc.
// Copyright © 2023-2024 Apple Inc.
#include <sstream>
#include <pybind11/pybind11.h>
#include <nanobind/nanobind.h>
#include <nanobind/stl/string.h>
#include "mlx/device.h"
#include "mlx/utils.h"
namespace py = pybind11;
using namespace py::literals;
namespace nb = nanobind;
using namespace nb::literals;
using namespace mlx::core;
void init_device(py::module_& m) {
auto device_class = py::class_<Device>(
void init_device(nb::module_& m) {
auto device_class = nb::class_<Device>(
m, "Device", R"pbdoc(A device to run operations on.)pbdoc");
py::enum_<Device::DeviceType>(m, "DeviceType")
nb::enum_<Device::DeviceType>(m, "DeviceType")
.value("cpu", Device::DeviceType::cpu)
.value("gpu", Device::DeviceType::gpu)
.export_values()
.def(
"__eq__",
[](const Device::DeviceType& d1, const Device& d2) {
return d1 == d2;
},
py::prepend());
.def("__eq__", [](const Device::DeviceType& d, const nb::object& other) {
if (!nb::isinstance<Device>(other) &&
!nb::isinstance<Device::DeviceType>(other)) {
return false;
}
return d == nb::cast<Device>(other);
});
device_class.def(py::init<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
.def_readonly("type", &Device::type)
device_class.def(nb::init<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
.def_ro("type", &Device::type)
.def(
"__repr__",
[](const Device& d) {
@@ -34,11 +36,15 @@ void init_device(py::module_& m) {
os << d;
return os.str();
})
.def("__eq__", [](const Device& d1, const Device& d2) {
return d1 == d2;
.def("__eq__", [](const Device& d, const nb::object& other) {
if (!nb::isinstance<Device>(other) &&
!nb::isinstance<Device::DeviceType>(other)) {
return false;
}
return d == nb::cast<Device>(other);
});
py::implicitly_convertible<Device::DeviceType, Device>();
nb::implicitly_convertible<Device::DeviceType, Device>();
m.def(
"default_device",