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

Fix stub generation, change graph exporting for arrows to go to outputs (#455)

Browse Source
This commit is contained in:
Awni Hannun
2024-01-14 14:06:16 -08:00
committed by GitHub
parent 6e81c3e164
commit 41cc7bdfdb
5 changed files with 49 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/device.cpp
+2 -2
View File
@@ -12,6 +12,7 @@ using namespace py::literals;
using namespace mlx::core;
void init_device(py::module_& m) {
auto device_class = py::class_<Device>(m, "Device");
py::enum_<Device::DeviceType>(m, "DeviceType")
.value("cpu", Device::DeviceType::cpu)
.value("gpu", Device::DeviceType::gpu)
@@ -23,8 +24,7 @@ void init_device(py::module_& m) {
},
py::prepend());
py::class_<Device>(m, "Device")
.def(py::init<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
device_class.def(py::init<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
.def_readonly("type", &Device::type)
.def(
"__repr__",