Move to nanobind v2 (#1316)

Author: awni

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.24
-mlx>=0.9.0
-nanobind@git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
+mlx>=0.16.2
+nanobind==2.0

python/src/array.cpp

@@ -126,7 +126,7 @@ void init_array(nb::module_& m) {
   m.attr("float32") = nb::cast(float32);
   m.attr("bfloat16") = nb::cast(bfloat16);
   m.attr("complex64") = nb::cast(complex64);
-  nb::class_<Dtype::Category>(
+  nb::enum_<Dtype::Category>(
       m,
       "DtypeCategory",
       R"pbdoc(
@@ -165,16 +165,16 @@ void init_array(nb::module_& m) {
               * :ref:`complex64 <data_types>`
 
       See also :func:`~mlx.core.issubdtype`.
-      )pbdoc");
-  m.attr("complexfloating") = nb::cast(complexfloating);
-  m.attr("floating") = nb::cast(floating);
-  m.attr("inexact") = nb::cast(inexact);
-  m.attr("signedinteger") = nb::cast(signedinteger);
-  m.attr("unsignedinteger") = nb::cast(unsignedinteger);
-  m.attr("integer") = nb::cast(integer);
-  m.attr("number") = nb::cast(number);
-  m.attr("generic") = nb::cast(generic);
-
+      )pbdoc")
+      .value("complexfloating", complexfloating)
+      .value("floating", floating)
+      .value("inexact", inexact)
+      .value("signedinteger", signedinteger)
+      .value("unsignedinteger", unsignedinteger)
+      .value("integer", integer)
+      .value("number", number)
+      .value("generic", generic)
+      .export_values();
   nb::class_<ArrayAt>(
       m,
       "_ArrayAt",

setup.py

@@ -176,7 +176,7 @@ def run(self) -> None:
         include_package_data=True,
         extras_require={
             "dev": [
-                "nanobind@git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4",
+                "nanobind==2.0",
                 "numpy",
                 "pre-commit",
                 "setuptools>=42",