| |
|
|
| #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"); |
| } |
|
|
