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 <[email protected]>
Co-authored-by: Awni Hannun <[email protected]>

Author: 3 people

docs/src/usage/numpy.rst

@@ -3,7 +3,11 @@
 Conversion to NumPy and Other Frameworks
 ========================================
 
-MLX array implements the `Python Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_.
+MLX array supports conversion between other frameworks with either:  
+
+* The `Python Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_. 
+* `DLPack <https://dmlc.github.io/dlpack/latest/>`_.  
+
 Let's convert an array to NumPy and back.
 
 .. code-block:: python

python/src/array.cpp

@@ -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__",

python/tests/test_array.py

@@ -161,6 +161,19 @@ def test_list_not_equals_array(self):
         self.assertTrue(a != b)
         self.assertTrue(a != c)
 
+    def test_dlx_device_type(self):
+        a = mx.array([1, 2, 3])
+        device_type, device_id = a.__dlpack_device__()
+        self.assertIn(device_type, [1, 8])
+        self.assertEqual(device_id, 0)
+
+        if device_type == 8:
+            # Additional check if Metal is supposed to be available
+            self.assertTrue(mx.metal.is_available())
+        elif device_type == 1:
+            # Additional check if CPU is the fallback
+            self.assertFalse(mx.metal.is_available())
+
     def test_tuple_not_equals_array(self):
         a = mx.array([1, 2, 3])
         b = (1, 2, 3)