|
|
|
|
|
|
|
|
#include <sstream> |
|
|
|
|
|
#include <nanobind/nanobind.h> |
|
|
#include <nanobind/stl/string.h> |
|
|
|
|
|
#include "mlx/device.h" |
|
|
#include "mlx/utils.h" |
|
|
|
|
|
namespace mx = mlx::core; |
|
|
namespace nb = nanobind; |
|
|
using namespace nb::literals; |
|
|
|
|
|
void init_device(nb::module_& m) { |
|
|
auto device_class = nb::class_<mx::Device>( |
|
|
m, "Device", R"pbdoc(A device to run operations on.)pbdoc"); |
|
|
nb::enum_<mx::Device::DeviceType>(m, "DeviceType") |
|
|
.value("cpu", mx::Device::DeviceType::cpu) |
|
|
.value("gpu", mx::Device::DeviceType::gpu) |
|
|
.export_values() |
|
|
.def( |
|
|
"__eq__", |
|
|
[](const mx::Device::DeviceType& d, const nb::object& other) { |
|
|
if (!nb::isinstance<mx::Device>(other) && |
|
|
!nb::isinstance<mx::Device::DeviceType>(other)) { |
|
|
return false; |
|
|
} |
|
|
return d == nb::cast<mx::Device>(other); |
|
|
}); |
|
|
|
|
|
device_class |
|
|
.def(nb::init<mx::Device::DeviceType, int>(), "type"_a, "index"_a = 0) |
|
|
.def_ro("type", &mx::Device::type) |
|
|
.def( |
|
|
"__repr__", |
|
|
[](const mx::Device& d) { |
|
|
std::ostringstream os; |
|
|
os << d; |
|
|
return os.str(); |
|
|
}) |
|
|
.def("__eq__", [](const mx::Device& d, const nb::object& other) { |
|
|
if (!nb::isinstance<mx::Device>(other) && |
|
|
!nb::isinstance<mx::Device::DeviceType>(other)) { |
|
|
return false; |
|
|
} |
|
|
return d == nb::cast<mx::Device>(other); |
|
|
}); |
|
|
|
|
|
nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>(); |
|
|
|
|
|
m.def( |
|
|
"default_device", |
|
|
&mx::default_device, |
|
|
R"pbdoc(Get the default device.)pbdoc"); |
|
|
m.def( |
|
|
"set_default_device", |
|
|
&mx::set_default_device, |
|
|
"device"_a, |
|
|
R"pbdoc(Set the default device.)pbdoc"); |
|
|
m.def( |
|
|
"is_available", |
|
|
&mx::is_available, |
|
|
"device"_a, |
|
|
R"pbdoc(Check if a back-end is available for the given device.)pbdoc"); |
|
|
} |
|
|
|
