Bump nanobind to 2.4 + fix (#1710)

* bump nanobind to 2.4 + fix

* fix

python/src/array.cpp

@@ -193,6 +193,15 @@ void init_array(nb::module_& m) {
       .def("maximum", &ArrayAt::maximum, "value"_a)
       .def("minimum", &ArrayAt::minimum, "value"_a);
 
+  nb::class_<ArrayLike>(
+      m,
+      "ArrayLike",
+      R"pbdoc(
+        Any Python object which has an ``__mlx__array__`` method that
+        returns an :obj:`array`.
+      )pbdoc")
+      .def(nb::init_implicit<nb::object>());
+
   nb::class_<ArrayPythonIterator>(
       m,
       "ArrayIterator",

python/src/convert.cpp

@@ -477,7 +477,7 @@ mx::array create_array(ArrayInitType v, std::optional<mx::Dtype> t) {
   } else if (auto pv = std::get_if<mx::array>(&v); pv) {
     return mx::astype(*pv, t.value_or((*pv).dtype()));
   } else {
-    auto arr = to_array_with_accessor(std::get<nb::object>(v));
+    auto arr = to_array_with_accessor(std::get<ArrayLike>(v).obj);
     return mx::astype(arr, t.value_or(arr.dtype()));
   }
 }

python/src/convert.h

@@ -12,6 +12,11 @@
 namespace mx = mlx::core;
 namespace nb = nanobind;
 
+struct ArrayLike {
+  ArrayLike(nb::object obj) : obj(obj) {};
+  nb::object obj;
+};
+
 using ArrayInitType = std::variant<
     nb::bool_,
     nb::int_,
@@ -23,7 +28,7 @@ using ArrayInitType = std::variant<
     std::complex<float>,
     nb::list,
     nb::tuple,
-    nb::object>;
+    ArrayLike>;
 
 mx::array nd_array_to_mlx(
     nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,

python/src/utils.cpp

@@ -28,7 +28,7 @@ mx::array to_array(
              pv) {
     return nd_array_to_mlx(*pv, dtype);
   } else {
-    return to_array_with_accessor(std::get<nb::object>(v));
+    return to_array_with_accessor(std::get<ArrayLike>(v).obj);
   }
 }
 
@@ -42,14 +42,15 @@ std::pair<mx::array, mx::array> to_arrays(
   // - If neither is an array convert to arrays but leave their types alone
   auto is_mlx_array = [](const ScalarOrArray& x) {
     return std::holds_alternative<mx::array>(x) ||
-        std::holds_alternative<nb::object>(x) &&
-        nb::hasattr(std::get<nb::object>(x), "__mlx_array__");
+        std::holds_alternative<ArrayLike>(x) &&
+        nb::hasattr(std::get<ArrayLike>(x).obj, "__mlx_array__");
   };
   auto get_mlx_array = [](const ScalarOrArray& x) {
     if (auto px = std::get_if<mx::array>(&x); px) {
       return *px;
     } else {
-      return nb::cast<mx::array>(std::get<nb::object>(x).attr("__mlx_array__"));
+      return nb::cast<mx::array>(
+          std::get<ArrayLike>(x).obj.attr("__mlx_array__"));
     }
   };
 

python/src/utils.h

@@ -11,6 +11,7 @@
 #include <nanobind/stl/variant.h>
 
 #include "mlx/array.h"
+#include "python/src/convert.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;
@@ -25,7 +26,7 @@ using ScalarOrArray = std::variant<
     // Must be above complex
     nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>,
     std::complex<float>,
-    nb::object>;
+    ArrayLike>;
 
 inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
   std::vector<int> axes;
@@ -43,8 +44,9 @@ inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
 inline bool is_comparable_with_array(const ScalarOrArray& v) {
   // Checks if the value can be compared to an array (or is already an
   // mlx array)
-  if (auto pv = std::get_if<nb::object>(&v); pv) {
-    return nb::isinstance<mx::array>(*pv) || nb::hasattr(*pv, "__mlx_array__");
+  if (auto pv = std::get_if<ArrayLike>(&v); pv) {
+    auto obj = (*pv).obj;
+    return nb::isinstance<mx::array>(obj) || nb::hasattr(obj, "__mlx_array__");
   } else {
     // If it's not an object, it's a scalar (nb::int_, nb::float_, etc.)
     // and can be compared to an array
@@ -53,7 +55,7 @@ inline bool is_comparable_with_array(const ScalarOrArray& v) {
 }
 
 inline nb::handle get_handle_of_object(const ScalarOrArray& v) {
-  return std::get<nb::object>(v).ptr();
+  return std::get<ArrayLike>(v).obj.ptr();
 }
 
 inline void throw_invalid_operation(