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

Remove "using namespace mlx::core" in python/src (#1689)

Browse Source
This commit is contained in:
Cheng
2024-12-12 08:45:39 +09:00
committed by GitHub
parent f3dfa36a3a
commit 0bf19037ca
22 changed files with 1423 additions and 1302 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/device.cpp
+25 -22
View File
@@ -8,51 +8,54 @@
#include "mlx/device.h"
#include "mlx/utils.h"
namespace mx = mlx::core;
namespace nb = nanobind;
using namespace nb::literals;
using namespace mlx::core;
void init_device(nb::module_& m) {
auto device_class = nb::class_<Device>(
auto device_class = nb::class_<mx::Device>(
m, "Device", R"pbdoc(A device to run operations on.)pbdoc");
nb::enum_<Device::DeviceType>(m, "DeviceType")
.value("cpu", Device::DeviceType::cpu)
.value("gpu", Device::DeviceType::gpu)
nb::enum_<mx::Device::DeviceType>(m, "DeviceType")
.value("cpu", mx::Device::DeviceType::cpu)
.value("gpu", mx::Device::DeviceType::gpu)
.export_values()
.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);
});
.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<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
.def_ro("type", &Device::type)
device_class
.def(nb::init<mx::Device::DeviceType, int>(), "type"_a, "index"_a = 0)
.def_ro("type", &mx::Device::type)
.def(
"__repr__",
[](const Device& d) {
[](const mx::Device& d) {
std::ostringstream os;
os << d;
return os.str();
})
.def("__eq__", [](const Device& d, const nb::object& other) {
if (!nb::isinstance<Device>(other) &&
!nb::isinstance<Device::DeviceType>(other)) {
.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<Device>(other);
return d == nb::cast<mx::Device>(other);
});
nb::implicitly_convertible<Device::DeviceType, Device>();
nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>();
m.def(
"default_device",
&default_device,
&mx::default_device,
R"pbdoc(Get the default device.)pbdoc");
m.def(
"set_default_device",
&set_default_device,
&mx::set_default_device,
"device"_a,
R"pbdoc(Set the default device.)pbdoc");
}