mirror of
https://github.com/ml-explore/mlx.git
synced 2025-11-01 16:58:08 +08:00
feat: Added dlpack device (#1165)
* feat: Added dlpack device * feat: Added device_id to dlpack device * feat: Added device_id to dlpack device * doc: updated conversion docs * doc: updated numpy.rst dlpack information * doc: updated numpy.rst dlpack information * Update docs/src/usage/numpy.rst * Update docs/src/usage/numpy.rst --------- Co-authored-by: Venkat Ramnan Kalyanakumar <venkatramnankalyanakumar@Venkats-MacBook-Air.local> Co-authored-by: Awni Hannun <awni.hannun@gmail.com>
This commit is contained in:
@@ -10,11 +10,13 @@
|
||||
#include <nanobind/stl/variant.h>
|
||||
#include <nanobind/stl/vector.h>
|
||||
|
||||
#include "mlx/backend/metal/metal.h"
|
||||
#include "python/src/buffer.h"
|
||||
#include "python/src/convert.h"
|
||||
#include "python/src/indexing.h"
|
||||
#include "python/src/utils.h"
|
||||
|
||||
#include "mlx/device.h"
|
||||
#include "mlx/ops.h"
|
||||
#include "mlx/transforms.h"
|
||||
#include "mlx/utils.h"
|
||||
@@ -353,6 +355,17 @@ void init_array(nb::module_& m) {
|
||||
new (&arr) array(nd_array_to_mlx(state, std::nullopt));
|
||||
})
|
||||
.def("__dlpack__", [](const array& a) { return mlx_to_dlpack(a); })
|
||||
.def(
|
||||
"__dlpack_device__",
|
||||
[](const array& a) {
|
||||
if (metal::is_available()) {
|
||||
// Metal device is available
|
||||
return nb::make_tuple(8, 0);
|
||||
} else {
|
||||
// CPU device
|
||||
return nb::make_tuple(1, 0);
|
||||
}
|
||||
})
|
||||
.def("__copy__", [](const array& self) { return array(self); })
|
||||
.def(
|
||||
"__deepcopy__",
|
||||
|
||||
Reference in New Issue
Block a user