python313Packages.mlx: 0.21.1 -> 0.26.3 (#410209)

maintainers/maintainer-list.nix

@@ -4108,6 +4108,12 @@
     name = "Cameron Smith";
     keys = [ { fingerprint = "3F14 C258 856E 88AE E0F9  661E FF04 3B36 8811 DD1C"; } ];
   };
+  cameronyule = {
+    email = "cameron@cameronyule.com";
+    github = "cameronyule";
+    githubId = 5451;
+    name = "Cameron Yule";
+  };
   camillemndn = {
     email = "camillemondon@free.fr";
     github = "camillemndn";

pkgs/development/python-modules/mlx/darwin-build-fixes.patch

@@ -0,0 +1,15 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index ab8aea44..9e1b06f4 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -100,9 +100,7 @@ elseif(MLX_BUILD_METAL)
+   endif()
+ 
+   # Throw an error if xcrun not found
+-  execute_process(
+-    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+-    OUTPUT_VARIABLE MACOS_SDK_VERSION COMMAND_ERROR_IS_FATAL ANY)
++  set(MACOS_SDK_VERSION @sdkVersion@)
+ 
+   if(${MACOS_SDK_VERSION} LESS 14.0)
+     message(

pkgs/development/python-modules/mlx/default.nix

@@ -1,14 +1,28 @@
 {
-  lib,
-  fetchFromGitHub,
   buildPythonPackage,
-  pybind11,
+  fetchFromGitHub,
-  cmake,
+  lib,
-  xcbuild,
+  replaceVars,
-  zsh,
+  stdenv,
-  blas,
+
-  lapack,
+  # build-system
   setuptools,
+
+  # nativeBuildInputs
+  cmake,
+
+  # buildInputs
+  apple-sdk_14,
+  fmt_10,
+  nanobind,
+  nlohmann_json,
+  pybind11,
+
+  # tests
+  numpy,
+  pytestCheckHook,
+  python,
+  runCommand,
 }:
 
 let
@@ -16,79 +30,125 @@ let
   gguf-tools = fetchFromGitHub {
     owner = "antirez";
     repo = "gguf-tools";
-    rev = "af7d88d808a7608a33723fba067036202910acb3";
+    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
-    hash = "sha256-LqNvnUbmq0iziD9VP5OTJCSIy+y/hp5lKCUV7RtKTvM=";
+    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
   };
-  nlohmann_json = fetchFromGitHub {
+
-    owner = "nlohmann";
+  mlx = buildPythonPackage rec {
-    repo = "json";
-    rev = "v3.11.3";
-    hash = "sha256-7F0Jon+1oWL7uqet5i1IgHX0fUw/+z0QwEcA3zs5xHg=";
-  };
-in
-buildPythonPackage rec {
     pname = "mlx";
-  version = "0.21.1";
+    version = "0.26.3";
+    pyproject = true;
 
     src = fetchFromGitHub {
       owner = "ml-explore";
       repo = "mlx";
-    rev = "refs/tags/v${version}";
+      tag = "v${version}";
-    hash = "sha256-wxv9bA9e8VyFv/FMh63sUTTNgkXHGQJNQhLuVynczZA=";
+      hash = "sha256-hbqV/2KYGJ1gyExZd5bgaxTdhl5+Gext+U/+1KAztMU=";
     };
 
-  pyproject = true;
-
     patches = [
-    # With Darwin SDK 11 we cannot include vecLib/cblas_new.h, this needs to wait for PR #229210
+      (replaceVars ./darwin-build-fixes.patch {
-    # In the meantime, pretend Accelerate is not available and use blas/lapack instead.
+        sdkVersion = apple-sdk_14.version;
-    ./disable-accelerate.patch
+      })
     ];
 
     postPatch = ''
-    substituteInPlace CMakeLists.txt \
+      substituteInPlace pyproject.toml \
-      --replace "/usr/bin/xcrun" "${xcbuild}/bin/xcrun" \
+        --replace-fail "nanobind==2.4.0" "nanobind>=2.4.0"
+
+      substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+        --replace-fail "g++" "$CXX"
     '';
 
     dontUseCmakeConfigure = true;
 
+    enableParallelBuilding = true;
+
+    # Allows multiple cores to be used in Python builds.
+    postUnpack = ''
+      export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
+    '';
+
     # updates the wrong fetcher rev attribute
     passthru.skipBulkUpdate = true;
 
     env = {
       PYPI_RELEASE = version;
-    # we can't use Metal compilation with Darwin SDK 11
       CMAKE_ARGS = toString [
+        # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
+        # To build mlx with Metal support in Nix, you'd need to use one of the sandbox escape
+        # hatches which let you interact with a native install of Xcode, such as `composeXcodeWrapper`
+        # or by changing the upstream (e.g., https://github.com/zed-industries/zed/discussions/7016).
         (lib.cmakeBool "MLX_BUILD_METAL" false)
         (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
-      (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_FMT" "${fmt_10.src}")
       ];
     };
 
-  nativeBuildInputs = [
+    build-system = [
-    cmake
-    pybind11
-    xcbuild
-    zsh
-    gguf-tools
-    nlohmann_json
       setuptools
     ];
 
-  buildInputs = [
+    nativeBuildInputs = [
-    blas
+      cmake
-    lapack
     ];
 
-  meta = with lib; {
+    buildInputs = [
+      apple-sdk_14
+      fmt_10
+      gguf-tools
+      nanobind
+      nlohmann_json
+      pybind11
+    ];
+
+    pythonImportsCheck = [ "mlx" ];
+
+    # Run the mlx Python test suite.
+    nativeCheckInputs = [
+      numpy
+      pytestCheckHook
+    ];
+
+    pytestFlagsArray = [
+      "python/tests/"
+    ];
+
+    # Additional testing by executing the example Python scripts supplied with mlx
+    # using the version of the library we've built.
+    passthru.tests = {
+      mlxTest =
+        runCommand "run-mlx-examples"
+          {
+            buildInputs = [ mlx ];
+            nativeBuildInputs = [ python ];
+          }
+          ''
+            cp ${src}/examples/python/logistic_regression.py .
+            ${python.interpreter} logistic_regression.py
+            rm logistic_regression.py
+
+            cp ${src}/examples/python/linear_regression.py .
+            ${python.interpreter} linear_regression.py
+            rm linear_regression.py
+
+            touch $out
+          '';
+    };
+
+    meta = {
       homepage = "https://github.com/ml-explore/mlx";
       description = "Array framework for Apple silicon";
-    changelog = "https://github.com/ml-explore/mlx/releases/tag/v${version}";
+      changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
-    license = licenses.mit;
+      license = lib.licenses.mit;
       platforms = [ "aarch64-darwin" ];
-    maintainers = with maintainers; [
+      maintainers = with lib.maintainers; [
         viraptor
         Gabriella439
+        cameronyule
       ];
     };
-}
+  };
+in
+mlx

pkgs/development/python-modules/mlx/disable-accelerate.patch

@@ -1,13 +0,0 @@
-diff --git a/CMakeLists.txt b/CMakeLists.txt
-index 2d6bef9..d099673 100644
---- a/CMakeLists.txt
-+++ b/CMakeLists.txt
-@@ -104,7 +104,7 @@ elseif (MLX_BUILD_METAL)
-     ${QUARTZ_LIB})
- endif()
- 
--find_library(ACCELERATE_LIBRARY Accelerate)
-+#find_library(ACCELERATE_LIBRARY Accelerate)
- if (MLX_BUILD_ARM AND ACCELERATE_LIBRARY)
-   message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
-   set(MLX_BUILD_ACCELERATE ON)