| 17 | using namespace nb::literals; |
| 18 | |
| 19 | void init_device(nb::module_& m) { |
| 20 | auto device_class = nb::class_<mx::Device>( |
| 21 | m, "Device", R"pbdoc(A device to run operations on.)pbdoc"); |
| 22 | nb::enum_<mx::Device::DeviceType>(m, "DeviceType") |
| 23 | .value("cpu", mx::Device::DeviceType::cpu) |
| 24 | .value("gpu", mx::Device::DeviceType::gpu) |
| 25 | .export_values() |
| 26 | .def( |
| 27 | "__eq__", |
| 28 | [](const mx::Device::DeviceType& d, const nb::object& other) { |
| 29 | if (!nb::isinstance<mx::Device>(other) && |
| 30 | !nb::isinstance<mx::Device::DeviceType>(other)) { |
| 31 | return false; |
| 32 | } |
| 33 | return d == nb::cast<mx::Device>(other); |
| 34 | }); |
| 35 | |
| 36 | device_class |
| 37 | .def(nb::init<mx::Device::DeviceType, int>(), "type"_a, "index"_a = 0) |
| 38 | .def_ro("type", &mx::Device::type) |
| 39 | .def( |
| 40 | "__repr__", |
| 41 | [](const mx::Device& d) { |
| 42 | std::ostringstream os; |
| 43 | os << d; |
| 44 | return os.str(); |
| 45 | }) |
| 46 | .def("__eq__", [](const mx::Device& d, const nb::object& other) { |
| 47 | if (!nb::isinstance<mx::Device>(other) && |
| 48 | !nb::isinstance<mx::Device::DeviceType>(other)) { |
| 49 | return false; |
| 50 | } |
| 51 | return d == nb::cast<mx::Device>(other); |
| 52 | }); |
| 53 | |
| 54 | nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>(); |
| 55 | |
| 56 | m.def( |
| 57 | "default_device", |
| 58 | &mx::default_device, |
| 59 | R"pbdoc(Get the default device.)pbdoc"); |
| 60 | m.def( |
| 61 | "set_default_device", |
| 62 | &mx::set_default_device, |
| 63 | "device"_a, |
| 64 | R"pbdoc(Set the default device.)pbdoc"); |
| 65 | m.def( |
| 66 | "is_available", |
| 67 | &mx::is_available, |
| 68 | "device"_a, |
| 69 | R"pbdoc(Check if a back-end is available for the given device.)pbdoc"); |
| 70 | m.def( |
| 71 | "device_count", |
| 72 | &mx::device_count, |
| 73 | "device_type"_a, |
| 74 | R"pbdoc( |
| 75 | Get the number of available devices for the given device type. |
| 76 | |