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

@@ -1,9 +1,4 @@
 // Copyright © 2023-2024 Apple Inc.
-#include <cstdint>
-#include <cstring>
-#include <sstream>
-#include <tuple>
-
 #include <nanobind/ndarray.h>
 #include <nanobind/stl/complex.h>
 #include <nanobind/stl/optional.h>
@@ -13,6 +8,11 @@
 #include <nanobind/stl/vector.h>
 #include <nanobind/typing.h>
 
+#include <cstdint>
+#include <cstring>
+#include <sstream>
+#include <tuple>
+
 #include "mlx/backend/metal/metal.h"
 #include "mlx/utils.h"
 #include "python/src/buffer.h"
@@ -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[] = {
@@ -1586,5 +1596,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/convert.cpp

@@ -1,13 +1,13 @@
 // Copyright © 2024 Apple Inc.
 
+#include <nanobind/stl/complex.h>
+#include <nanobind/stl/string.h>
+
 #include <algorithm>
 #include <limits>
 #include <sstream>
 #include <tuple>
 
-#include <nanobind/stl/complex.h>
-#include <nanobind/stl/string.h>
-
 #include "python/src/convert.h"
 #include "python/src/utils.h"
 

python/src/device.cpp

@@ -1,14 +1,14 @@
 // Copyright © 2023-2025 Apple Inc.
 
-#include <optional>
-#include <sstream>
-
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/unordered_map.h>
 #include <nanobind/stl/variant.h>
 
+#include <optional>
+#include <sstream>
+
 #include "mlx/device.h"
 #include "mlx/utils.h"
 
@@ -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/indexing.cpp

@@ -1,10 +1,10 @@
 // Copyright © 2023-2024 Apple Inc.
+#include <nanobind/ndarray.h>
+
 #include <numeric>
 #include <optional>
 #include <sstream>
 
-#include <nanobind/ndarray.h>
-
 #include "mlx/dtype.h"
 #include "mlx/ops.h"
 #include "python/src/convert.h"

python/src/linalg.cpp

@@ -1,13 +1,13 @@
 // Copyright © 2023-2024 Apple Inc.
 
-#include <variant>
-
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/pair.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
 
+#include <variant>
+
 #include "mlx/linalg.h"
 #include "python/src/small_vector.h"
 

python/src/memory.cpp

@@ -1,8 +1,8 @@
 // Copyright © 2025-2026 Apple Inc.
 
-#include "mlx/memory.h"
 #include <nanobind/nanobind.h>
 
+#include "mlx/memory.h"
 #include "python/src/trees.h"
 
 namespace mx = mlx::core;

python/src/metal.cpp

@@ -1,13 +1,13 @@
 // Copyright © 2023-2024 Apple Inc.
-#include <iostream>
-
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/unordered_map.h>
 #include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
 
+#include <iostream>
+
 #include "mlx/backend/metal/metal.h"
 #include "mlx/device.h"
 #include "mlx/memory.h"