Improve DLPack-compatible array imports (#3495)

Author: XXXXRT666

docs/src/usage/numpy.rst

@@ -75,22 +75,16 @@ even though no in-place operations on MLX memory are executed.
 PyTorch
 -------
 
-.. warning::
-
-   PyTorch Support for :obj:`memoryview` is experimental and can break for
-   multi-dimensional arrays. Casting to NumPy first is advised for now.
-
-PyTorch supports the buffer protocol, but it requires an explicit
-:obj:`memoryview`.
+PyTorch supports DLPack inputs and can import MLX arrays directly.
 
 .. code-block:: python
 
   import mlx.core as mx
   import torch
 
   a = mx.arange(3)
-  b = torch.tensor(memoryview(a))
-  c = mx.array(b)
+  b = torch.tensor(a)
+  c = mx.array(b.cpu())
 
 JAX
 ---

python/mlx/_stub_patterns.txt

@@ -1,12 +1,15 @@
 mlx.core.__prefix__:
-  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, ParamSpec, TypeVar
+  from typing import Any, Callable, Dict, List, Optional, Protocol, Sequence, Tuple, Union, ParamSpec, TypeVar
   import sys
   if sys.version_info >= (3, 10):
     from typing import TypeAlias
   else:
     from typing_extensions import TypeAlias
   P = ParamSpec("P")
   R = TypeVar("R")
+  class DLPackCompatible(Protocol):
+    __dlpack__: Callable[..., Any]
+    __dlpack_device__: Callable[..., Any]
 
 mlx.core.__suffix__:
   from typing import Union

python/src/array.cpp

@@ -300,7 +300,7 @@ void init_array(nb::module_& m) {
           "val"_a,
           "dtype"_a = nb::none(),
           nb::sig(
-              "def __init__(self: array, val: Union[scalar, list, tuple, numpy.ndarray, array], dtype: Optional[Dtype] = None)"))
+              "def __init__(self: array, val: Union[scalar, list, tuple, DLPackCompatible, array], dtype: Optional[Dtype] = None)"))
       .def_prop_ro(
           "size",
           &mx::array::size,
@@ -479,7 +479,7 @@ void init_array(nb::module_& m) {
               throw std::invalid_argument(
                   "Invalid pickle state: expected (ndarray, Dtype::Val)");
             }
-            using ND = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>;
+            using ND = nb::ndarray<nb::ro, nb::c_contig>;
             ND nd = nb::cast<ND>(state[0]);
             auto val = static_cast<mx::Dtype::Val>(nb::cast<uint8_t>(state[1]));
             if (val == mx::Dtype::Val::bfloat16) {

python/src/convert.cpp

@@ -33,7 +33,7 @@ int check_shape_dim(int64_t dim) {
 
 template <typename T>
 mx::array nd_array_to_mlx_contiguous(
-    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
+    nb::ndarray<nb::ro, nb::c_contig> nd_array,
     const mx::Shape& shape,
     mx::Dtype dtype) {
   // Make a copy of the numpy buffer
@@ -43,9 +43,14 @@ mx::array nd_array_to_mlx_contiguous(
 }
 
 mx::array nd_array_to_mlx(
-    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
+    nb::ndarray<nb::ro, nb::c_contig> nd_array,
     std::optional<mx::Dtype> dtype,
     std::optional<nb::dlpack::dtype> nb_dtype) {
+  if (nd_array.device_type() != nb::device::cpu::value) {
+    throw std::invalid_argument(
+        "Cannot convert non-CPU DLPack array to mlx array.");
+  }
+
   // Compute the shape and size
   mx::Shape shape;
   shape.reserve(nd_array.ndim());
@@ -495,7 +500,7 @@ mx::array create_array(nb::object v, std::optional<mx::Dtype> t) {
     auto arr = nb::cast<mx::array>(v);
     return mx::astype(arr, t.value_or(arr.dtype()));
   } else if (nb::ndarray_check(v)) {
-    using ContigArray = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>;
+    using ContigArray = nb::ndarray<nb::ro, nb::c_contig>;
     ContigArray nd;
     std::optional<nb::dlpack::dtype> nb_dtype;
     // Nanobind does not recognize bfloat16 numpy array:

python/src/convert.h

@@ -62,7 +62,7 @@ struct ArrayLike {
 };
 
 mx::array nd_array_to_mlx(
-    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
+    nb::ndarray<nb::ro, nb::c_contig> nd_array,
     std::optional<mx::Dtype> mx_dtype,
     std::optional<nb::dlpack::dtype> nb_dtype = std::nullopt);
 

python/src/indexing.cpp

@@ -926,7 +926,7 @@ mlx_compute_slice_update_args(
 }
 
 std::optional<mx::array> extract_boolean_mask(const nb::object& obj) {
-  using NDArray = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>;
+  using NDArray = nb::ndarray<nb::ro, nb::c_contig>;
   if (nb::isinstance<nb::bool_>(obj)) {
     return mx::array(nb::cast<bool>(obj), mx::bool_);
   } else if (nb::isinstance<mx::array>(obj)) {

python/src/utils.cpp

@@ -38,9 +38,7 @@ mx::array to_array(
     return mx::array(static_cast<mx::complex64_t>(*pv), mx::complex64);
   } else if (auto pv = std::get_if<mx::array>(&v); pv) {
     return *pv;
-  } else if (auto pv = std::get_if<
-                 nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>>(&v);
-             pv) {
+  } else if (auto pv = std::get_if<nb::ndarray<nb::ro, nb::c_contig>>(&v); pv) {
     return nd_array_to_mlx(*pv, dtype);
   } else {
     return to_array_with_accessor(std::get<ArrayLike>(v).obj);

python/src/utils.h

@@ -24,7 +24,7 @@ using ScalarOrArray = std::variant<
     // Must be above ndarray
     mx::array,
     // Must be above complex
-    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>,
+    nb::ndarray<nb::ro, nb::c_contig>,
     std::complex<float>,
     ArrayLike>;
 

python/tests/test_array.py

@@ -20,9 +20,17 @@
     import tensorflow as tf
 
     has_tf = True
-except ImportError as e:
+except ImportError:
     has_tf = False
 
+try:
+    import torch
+
+    has_torch_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+except ImportError:
+    torch = None
+    has_torch_mps = False
+
 
 class TestVersion(mlx_tests.MLXTestCase):
     def test_version(self):
@@ -2039,6 +2047,21 @@ def test_dlpack(self):
         y = np.from_dlpack(x)
         self.assertTrue(mx.array_equal(y, x))
 
+    @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required")
+    def test_torch_mps_dlpack_non_cpu_error(self):
+        x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4)
+        self.assertEqual(x.__dlpack_device__()[0], 8)
+
+        with self.assertRaisesRegex(ValueError, "non-CPU DLPack"):
+            mx.array(x)
+
+        a = mx.array([1])
+        b = torch.tensor([2])
+        self.assertTrue(mx.array_equal(a + b, mx.array([3])))
+
+        with self.assertRaisesRegex(ValueError, "non-CPU DLPack"):
+            a + b.to("mps")
+
     def test_getitem_with_list(self):
         a = mx.array([1, 2, 3, 4, 5])
         idx = [0, 2, 4]