Add Python 3.15 and 3.15t support

.github/workflows/release.yml

@@ -58,7 +58,7 @@ jobs:
       matrix:
         os: ['Linux', 'Windows']
         arch: ['x86_64', 'aarch64']
-        python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t']
+        python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t', '3.15', '3.15t']
         # There is no cp310 binary for Windows on arm.
         exclude:
           - os: 'Windows'

CMakeLists.txt

@@ -396,7 +396,7 @@ if(MLX_BUILD_PYTHON_BINDINGS)
   FetchContent_Declare(
     nanobind
     GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-    GIT_TAG v3.0.1
+    GIT_TAG v3.1.0
     GIT_SHALLOW TRUE
     EXCLUDE_FROM_ALL)
   FetchContent_MakeAvailable(nanobind)

docs/src/dev/extensions.rst

@@ -679,7 +679,7 @@ Finally, we build the nanobind_ bindings
 
     nanobind_add_module(
       _ext
-      NB_STATIC STABLE_ABI LTO NOMINSIZE
+      NB_STATIC FREE_THREADED LTO NOMINSIZE
       NB_DOMAIN mlx
       ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
     )

examples/extensions/CMakeLists.txt

@@ -23,7 +23,7 @@ find_package(
 FetchContent_Declare(
   nanobind
   GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-  GIT_TAG v3.0.1
+  GIT_TAG v3.1.0
   GIT_SHALLOW TRUE
   EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(nanobind)
@@ -73,7 +73,7 @@ endif()
 nanobind_add_module(
   _ext
   NB_STATIC
-  STABLE_ABI
+  FREE_THREADED
   LTO
   NOMINSIZE
   NB_DOMAIN

python/src/CMakeLists.txt

@@ -1,7 +1,6 @@
 nanobind_add_module(
   core
   NB_STATIC
-  STABLE_ABI
   FREE_THREADED
   LTO
   NOMINSIZE

python/src/array.cpp

@@ -129,9 +129,10 @@ void init_array(nb::module_& m) {
             return nb::isinstance<mx::Dtype>(other) &&
                 t == nb::cast<mx::Dtype>(other);
           })
-      .def("__hash__", [](const mx::Dtype& t) {
-        return static_cast<int64_t>(t.val());
-      });
+      .def(
+          "__hash__",
+          [](const mx::Dtype& t) { return static_cast<int64_t>(t.val()); })
+      .freeze();
 
   m.attr("bool_") = nb::cast(mx::bool_);
   m.attr("uint8") = nb::cast(mx::uint8);
@@ -229,13 +230,16 @@ void init_array(nb::module_& m) {
           &mx::finfo::smallest_normal,
           R"pbdoc(The smallest positive normal number.)pbdoc")
       .def_ro("dtype", &mx::finfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc")
-      .def("__repr__", [](const mx::finfo& f) {
-        std::ostringstream os;
-        os << "finfo("
-           << "min=" << f.min << ", max=" << f.max << ", dtype=" << f.dtype
-           << ")";
-        return os.str();
-      });
+      .def(
+          "__repr__",
+          [](const mx::finfo& f) {
+            std::ostringstream os;
+            os << "finfo("
+               << "min=" << f.min << ", max=" << f.max << ", dtype=" << f.dtype
+               << ")";
+            return os.str();
+          })
+      .freeze();
 
   nb::class_<mx::iinfo>(
       m,
@@ -253,13 +257,16 @@ void init_array(nb::module_& m) {
           &mx::iinfo::max,
           R"pbdoc(The largest representable number.)pbdoc")
       .def_ro("dtype", &mx::iinfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc")
-      .def("__repr__", [](const mx::iinfo& i) {
-        std::ostringstream os;
-        os << "iinfo("
-           << "min=" << i.min << ", max=" << i.max << ", dtype=" << i.dtype
-           << ")";
-        return os.str();
-      });
+      .def(
+          "__repr__",
+          [](const mx::iinfo& i) {
+            std::ostringstream os;
+            os << "iinfo("
+               << "min=" << i.min << ", max=" << i.max << ", dtype=" << i.dtype
+               << ")";
+            return os.str();
+          })
+      .freeze();
 
   nb::class_<ArrayAt>(
       m,
@@ -274,7 +281,8 @@ void init_array(nb::module_& m) {
       .def("multiply", &ArrayAt::multiply, "value"_a)
       .def("divide", &ArrayAt::divide, "value"_a)
       .def("maximum", &ArrayAt::maximum, "value"_a)
-      .def("minimum", &ArrayAt::minimum, "value"_a);
+      .def("minimum", &ArrayAt::minimum, "value"_a)
+      .freeze();
 
   nb::class_<ArrayLike>(
       m,
@@ -283,7 +291,8 @@ void init_array(nb::module_& m) {
         Any Python object which has an ``__mlx__array__`` method that
         returns an :obj:`array`.
       )pbdoc")
-      .def(nb::init_implicit<nb::object>());
+      .def(nb::init_implicit<nb::object>())
+      .freeze();
 
   nb::class_<ArrayPythonIterator>(
       m,
@@ -292,7 +301,8 @@ void init_array(nb::module_& m) {
       A helper object to iterate over the 1st dimension of an array.
       )pbdoc")
       .def("__next__", &ArrayPythonIterator::next)
-      .def("__iter__", [](const ArrayPythonIterator& it) { return it; });
+      .def("__iter__", [](const ArrayPythonIterator& it) { return it; })
+      .freeze();
 
   // Install buffer protocol functions
   PyType_Slot array_slots[] = {
@@ -1582,5 +1592,6 @@ void init_array(nb::module_& m) {
           "dtype"_a,
           nb::kw_only(),
           "stream"_a = nb::none(),
-          "See :func:`view`.");
+          "See :func:`view`.")
+      .freeze();
 }

python/src/device.cpp

@@ -43,13 +43,16 @@ void init_device(nb::module_& m) {
             os << d;
             return os.str();
           })
-      .def("__eq__", [](const mx::Device& d, const nb::object& other) {
-        if (!nb::isinstance<mx::Device>(other) &&
-            !nb::isinstance<mx::Device::DeviceType>(other)) {
-          return false;
-        }
-        return d == nb::cast<mx::Device>(other);
-      });
+      .def(
+          "__eq__",
+          [](const mx::Device& d, const nb::object& other) {
+            if (!nb::isinstance<mx::Device>(other) &&
+                !nb::isinstance<mx::Device::DeviceType>(other)) {
+              return false;
+            }
+            return d == nb::cast<mx::Device>(other);
+          })
+      .freeze();
 
   nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>();
 

python/src/distributed.cpp

@@ -51,7 +51,8 @@ void init_distributed(nb::module_& parent_module) {
               color (int): A value to group processes into subgroups.
               key (int, optional): A key to optionally change the rank ordering
                 of the processes.
-          )pbdoc");
+          )pbdoc")
+      .freeze();
 
   m.def(
       "is_available",

python/src/export.cpp

@@ -298,7 +298,8 @@ void init_export(nb::module_& m) {
             auto [args_, kwargs_] =
                 validate_and_extract_inputs(args, kwargs, "[export_function]");
             exporter(args_, kwargs_);
-          });
+          })
+      .freeze();
 
   m.def(
       "exporter",

python/src/mlx_func.cpp

@@ -2,17 +2,11 @@
 
 #include "python/src/mlx_func.h"
 
-#include <structmember.h>
-
-// A garbage collected function which wraps nb::cpp_function
-// See https://github.com/wjakob/nanobind/discussions/919
+// The wrapper owns callbacks because free-threaded nanobind functions are
+// immortal
 
 struct gc_func {
-  PyObject_HEAD
-      // Vector call implementation that forwards calls to nanobind
-      PyObject* (*vectorcall)(PyObject*, PyObject* const*, size_t, PyObject*);
-  // The nanobind wrapper func
-  PyObject* func;
+  std::unique_ptr<PyFunction> func;
 
   // The original wrapped func
   PyObject* orig_func;
@@ -22,97 +16,48 @@ struct gc_func {
 
 int gc_func_tp_traverse(PyObject* self, visitproc visit, void* arg) {
   Py_VISIT(Py_TYPE(self));
-  gc_func* w = (gc_func*)self;
-  Py_VISIT(w->func);
-  for (auto d : w->deps) {
-    Py_VISIT(d);
+  if (nb::inst_ready(self)) {
+    for (auto d : nb::inst_ptr<gc_func>(self)->deps) {
+      Py_VISIT(d);
+    }
   }
   return 0;
-};
+}
 
 int gc_func_tp_clear(PyObject* self) {
-  gc_func* w = (gc_func*)self;
-  Py_CLEAR(w->func);
+  auto* w = nb::inst_ptr<gc_func>(self);
+  w->orig_func = nullptr;
+  w->deps.clear();
+  w->func.reset();
   return 0;
 }
 
-PyObject* gc_func_get_doc(PyObject* self, void*) {
-  return PyObject_GetAttrString(((gc_func*)self)->func, "__doc__");
-}
-
-PyObject* gc_func_get_sig(PyObject* self, void*) {
-  return PyObject_GetAttrString(((gc_func*)self)->func, "__nb_signature__");
-}
-
-PyObject* gc_func_vectorcall(
-    PyObject* self,
-    PyObject* const* args,
-    size_t nargs,
-    PyObject* kwnames) {
-  return PyObject_Vectorcall(((gc_func*)self)->func, args, nargs, kwnames);
-}
-
-void gc_func_dealloc(PyObject* self) {
-  PyObject_GC_UnTrack(self);
-  Py_XDECREF(((gc_func*)self)->func);
-  PyObject_GC_Del(self);
-}
-
-static PyMemberDef gc_func_members[] = {
-    {"__vectorcalloffset__",
-     T_PYSSIZET,
-     (Py_ssize_t)offsetof(gc_func, vectorcall),
-     READONLY,
-     nullptr},
-    {nullptr, 0, 0, 0, nullptr}};
-
-static PyGetSetDef gc_func_getset[] = {
-    {"__doc__", gc_func_get_doc, nullptr, nullptr, nullptr},
-    {"__nb_signature__", gc_func_get_sig, nullptr, nullptr, nullptr},
-    {nullptr, nullptr, nullptr, nullptr, nullptr}};
-
 static PyObject* gc_func_getattro(PyObject* self, PyObject* name_) {
-  gc_func* w = (gc_func*)self;
-  return PyObject_GenericGetAttr(w->orig_func, name_);
+  return PyObject_GenericGetAttr(nb::inst_ptr<gc_func>(self)->orig_func, name_);
 }
 
 // Table of custom type slots we want to install
 PyType_Slot gc_func_slots[] = {
     {Py_tp_traverse, (void*)gc_func_tp_traverse},
     {Py_tp_clear, (void*)gc_func_tp_clear},
-    {Py_tp_getset, (void*)gc_func_getset},
     {Py_tp_getattro, (void*)gc_func_getattro},
-    {Py_tp_members, (void*)gc_func_members},
-    {Py_tp_call, (void*)PyVectorcall_Call},
-    {Py_tp_dealloc, (void*)gc_func_dealloc},
     {0, 0}};
 
-static PyType_Spec gc_func_spec = {
-    /* .name = */ "mlx.gc_func",
-    /* .basicsize = */ (int)sizeof(gc_func),
-    /* .itemsize = */ 0,
-    /* .flags = */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
-        Py_TPFLAGS_HAVE_VECTORCALL,
-    /* .slots = */ gc_func_slots};
-
-static PyTypeObject* gc_func_tp = nullptr;
-
 nb::callable mlx_func(
-    nb::object func,
+    std::unique_ptr<PyFunction> func,
     const nb::callable& orig_func,
     std::vector<PyObject*> deps) {
-  gc_func* r = (gc_func*)PyType_GenericAlloc(gc_func_tp, 0);
-  r->func = func.inc_ref().ptr();
-  r->orig_func = orig_func.ptr();
-  deps.push_back(r->orig_func);
-  r->deps = std::move(deps);
-  r->vectorcall = gc_func_vectorcall;
-  return nb::steal<nb::callable>((PyObject*)r);
+  deps.push_back(orig_func.ptr());
+  return nb::borrow<nb::callable>(
+      nb::cast(gc_func{std::move(func), orig_func.ptr(), std::move(deps)}));
 }
 
 void init_mlx_func(nb::module_& m) {
-  gc_func_tp = (PyTypeObject*)PyType_FromSpec(&gc_func_spec);
-  if (!gc_func_tp) {
-    nb::raise("Could not register MLX function type.");
-  }
+  nb::class_<gc_func>(m, "_gc_func", nb::type_slots(gc_func_slots))
+      .def(
+          "__call__",
+          [](gc_func& self, nb::args args, nb::kwargs kwargs) {
+            return (*self.func)(args, kwargs);
+          })
+      .freeze();
 }

python/src/mlx_func.h

@@ -2,30 +2,38 @@
 
 #pragma once
 
+#include <memory>
+#include <utility>
 #include <vector>
 
 #include <nanobind/nanobind.h>
-#include <nanobind/stl/function.h>
 
 namespace nb = nanobind;
 using namespace nb::literals;
 
+struct PyFunction {
+  virtual ~PyFunction() = default;
+  virtual nb::object operator()(nb::args& args, nb::kwargs& kwargs) = 0;
+};
+
 nb::callable mlx_func(
-    nb::object func,
+    std::unique_ptr<PyFunction> func,
     const nb::callable& orig_func,
     std::vector<PyObject*> deps);
 
 template <typename F, typename... Deps>
 nb::callable mlx_func(F func, const nb::callable& orig_func, Deps&&... deps) {
-  return mlx_func(
-      nb::cpp_function(std::move(func)),
-      orig_func,
-      std::vector<PyObject*>{deps.ptr()...});
-}
+  struct Callback : PyFunction {
+    F func;
+
+    explicit Callback(F func) : func(std::move(func)) {}
 
-template <typename... Deps>
-nb::callable
-mlx_func(nb::object func, const nb::callable& orig_func, Deps&&... deps) {
+    nb::object operator()(nb::args& args, nb::kwargs& kwargs) override {
+      return nb::cast(func(args, kwargs));
+    }
+  };
+  std::unique_ptr<PyFunction> callback =
+      std::make_unique<Callback>(std::move(func));
   return mlx_func(
-      std::move(func), orig_func, std::vector<PyObject*>{deps.ptr()...});
+      std::move(callback), orig_func, std::vector<PyObject*>{deps.ptr()...});
 }