Fetch nanobind with cmake (#2949)

.github/actions/build-linux-release/action.yml

@@ -18,12 +18,9 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Generate package stubs
+    - name: Build MLX
       shell: bash
-      run: |
-        pip install -e ".[dev]" -v
-        pip install typing_extensions
-        python setup.py generate_stubs
+      run: pip install -e . -v
 
     - name: Build Python package
       shell: bash

.github/actions/build-linux/action.yml

@@ -21,19 +21,13 @@ runs:
         if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
           # There is no GPU in arm64 runner, use a common arch.
           CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
-          # Can not build tests when the built executables can not run.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
+          # Can not build tests and stubs when the built executables can not run.
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
         fi
         pip install --no-build-isolation -e ".[dev]" -v
         # Pass the CMAKE_ARGS to following steps.
         echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT
 
-    - name: Generate package stubs
-      shell: sh
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
     - name: Build CPP only
       shell: bash
       run: |

.github/actions/build-macos/action.yml

@@ -11,15 +11,9 @@ runs:
       shell: bash -l {0}
       run: |
         pip install --upgrade pip
-        pip install cmake setuptools nanobind==2.10.2
+        pip install cmake setuptools typing_extensions
         pip install -e . -v
 
-    - name: Generate package stubs
-      shell: bash -l {0}
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
     - name: Install tests dependencies
       shell: bash -l {0}
       run: |

.github/actions/setup-linux/action.yml

@@ -48,9 +48,9 @@ runs:
         echo "::group::Setup Python venv"
         python -m venv .venv
         source .venv/bin/activate
-        pip install setuptools cmake nanobind==2.10.2
+        pip install setuptools cmake typing_extensions
         echo PATH=$PATH >> $GITHUB_ENV
-        # Make cmake search .venv for nanobind
+        # Search python packages in .venv
         echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
         echo "::endgroup::"
 

.github/workflows/release.yml

@@ -98,13 +98,8 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade pip
-          pip install cmake setuptools nanobind==2.10.2
+          pip install cmake setuptools typing_extensions
           pip install -e . -v
-      - name: Generate package stubs
-        shell: bash -l {0}
-        run: |
-          pip install typing_extensions
-          python setup.py generate_stubs
       - name: Build macOS 14 package
         uses: ./.github/actions/build-macos-release
         with:

CMakeLists.txt

@@ -41,6 +41,7 @@ option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
 option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
+option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
@@ -338,11 +339,13 @@ if(MLX_BUILD_PYTHON_BINDINGS)
     Python 3.10
     COMPONENTS Interpreter Development.Module
     REQUIRED)
-  execute_process(
-    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-    OUTPUT_VARIABLE nanobind_ROOT)
-  find_package(nanobind CONFIG REQUIRED)
+  FetchContent_Declare(
+    nanobind
+    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
+    GIT_TAG v2.10.2
+    GIT_SHALLOW TRUE
+    EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(nanobind)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()
 

docs/src/install.rst

@@ -128,13 +128,6 @@ Run the tests with:
 
   python -m unittest discover python/tests
 
-Optional: Install stubs to enable auto completions and type checking from your
-IDE:
-
-.. code-block:: shell
-
-  python setup.py generate_stubs
-
 C++ API
 ^^^^^^^
 

pyproject.toml

@@ -1,7 +1,7 @@
 [build-system]
 requires = [
   "setuptools>=80",
-  "nanobind==2.10.2",
   "cmake>=3.25",
+  "typing_extensions",
 ]
 build-backend = "setuptools.build_meta"

python/src/CMakeLists.txt

@@ -29,6 +29,37 @@ nanobind_add_module(
   ${CMAKE_CURRENT_SOURCE_DIR}/trees.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
 
+if(MLX_BUILD_PYTHON_STUBS)
+  nanobind_add_stub(
+    core_stub
+    # Run stubgen -m mlx.core -i python -p _stub_patterns.txt -o python/mlx
+    RECURSIVE
+    MODULE
+    "mlx.core"
+    PYTHON_PATH
+    "${CMAKE_CURRENT_SOURCE_DIR}/.."
+    PATTERN_FILE
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/_stub_patterns.txt"
+    OUTPUT_PATH
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx"
+    # Note that the list is passed to cmake for dependency managment and not
+    # used by stubgen.
+    OUTPUT
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/__init__.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/cuda.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/distributed.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/fast.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/fft.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/linalg.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/metal.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/random.pyi"
+    # Make this an optional installable component.
+    EXCLUDE_FROM_ALL
+    INSTALL_TIME
+    COMPONENT
+    core_stub)
+endif()
+
 if(NOT MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY)
   if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY)
     set(MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR})

setup.py

@@ -8,7 +8,7 @@ import subprocess
 from functools import partial
 from pathlib import Path
 
-from setuptools import Command, Extension, find_namespace_packages, setup
+from setuptools import Extension, find_namespace_packages, setup
 from setuptools.command.bdist_wheel import bdist_wheel
 from setuptools.command.build_ext import build_ext
 
@@ -155,43 +155,25 @@ class CMakeBuild(build_ext):
     def run(self):
         super().run()
 
+        ext = next(ext for ext in self.extensions if ext.name == "mlx.core")
+
         # Based on https://github.com/pypa/setuptools/blob/main/setuptools/command/build_ext.py#L102
         if self.inplace:
-            for ext in self.extensions:
-                if ext.name == "mlx.core":
-                    # Resolve inplace package dir
-                    build_py = self.get_finalized_command("build_py")
-                    inplace_file, regular_file = self._get_inplace_equivalent(
-                        build_py, ext
-                    )
+            # Resolve inplace package dir
+            build_py = self.get_finalized_command("build_py")
+            inplace_file, regular_file = self._get_inplace_equivalent(build_py, ext)
 
-                    inplace_dir = str(Path(inplace_file).parent.resolve())
-                    regular_dir = str(Path(regular_file).parent.resolve())
+            inplace_dir = str(Path(inplace_file).parent.resolve())
+            regular_dir = str(Path(regular_file).parent.resolve())
 
-                    self.copy_tree(regular_dir, inplace_dir)
+            self.copy_tree(regular_dir, inplace_dir)
 
-
-class GenerateStubs(Command):
-    user_options = []
-
-    def initialize_options(self):
-        pass
-
-    def finalize_options(self):
-        pass
-
-    def run(self) -> None:
-        out_path = "python/mlx"
-        stub_cmd = [
-            "python",
-            "-m",
-            "nanobind.stubgen",
-            "-m",
-            "mlx.core",
-            "-p",
-            "python/mlx/_stub_patterns.txt",
-        ]
-        subprocess.run(stub_cmd + ["-r", "-O", out_path])
+        # Build type stubs.
+        build_temp = Path(self.build_temp) / ext.name
+        subprocess.run(
+            ["cmake", "--install", build_temp, "--component", "core_stub"],
+            check=True,
+        )
 
 
 class MLXBdistWheel(bdist_wheel):
@@ -243,7 +225,6 @@ if __name__ == "__main__":
         ext_modules=[CMakeExtension("mlx.core")],
         cmdclass={
             "build_ext": CMakeBuild,
-            "generate_stubs": GenerateStubs,
             "bdist_wheel": MLXBdistWheel,
         },
     )