Merge pull request #286243 from viraptor/init-mlx

python3Packages.mlx: init at 0.1.0

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

@@ -0,0 +1,78 @@
+{ stdenv
+, lib
+, fetchFromGitHub
+, buildPythonPackage
+, python3Packages
+, pybind11
+, cmake
+, xcbuild
+, zsh
+, darwin
+, blas
+, lapack
+}:
+
+let
+  # static dependencies included directly during compilation
+  gguf-tools = fetchFromGitHub {
+    owner = "antirez";
+    repo = "gguf-tools";
+    rev = "af7d88d808a7608a33723fba067036202910acb3";
+    hash = "sha256-LqNvnUbmq0iziD9VP5OTJCSIy+y/hp5lKCUV7RtKTvM=";
+  };
+  nlohmann_json = fetchFromGitHub {
+    owner = "nlohmann";
+    repo = "json";
+    rev = "v3.11.3";
+    hash = "sha256-7F0Jon+1oWL7uqet5i1IgHX0fUw/+z0QwEcA3zs5xHg=";
+  };
+in
+buildPythonPackage rec {
+  pname = "mlx";
+  version = "0.1.0";
+
+  src = fetchFromGitHub {
+    owner = "ml-explore";
+    repo = "mlx";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-xNJPG8XGbC0fy6RGcn1cxCsejyHsgnV35PuP8F1I4R4=";
+  };
+
+  pyproject = true;
+
+  patches = [
+    # With Darwin SDK 11 we cannot include vecLib/cblas_new.h, this needs to wait for PR #229210
+    # In the meantime, pretend Accelerate is not available and use blas/lapack instead.
+    ./disable-accelerate.patch
+  ];
+
+  postPatch = ''
+    substituteInPlace CMakeLists.txt \
+      --replace "/usr/bin/xcrun" "${xcbuild}/bin/xcrun" \
+  '';
+
+  dontUseCmakeConfigure = true;
+
+  env = {
+    PYPI_RELEASE = version;
+    # we can't use Metal compilation with Darwin SDK 11
+    CMAKE_ARGS = toString [
+      (lib.cmakeBool "MLX_BUILD_METAL" false)
+      (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+      (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json}")
+    ];
+  };
+
+  nativeBuildInputs = [ cmake pybind11 xcbuild zsh gguf-tools nlohmann_json ] ++ (with python3Packages; [ setuptools ]);
+
+  buildInputs = [ blas lapack ];
+
+  meta = with lib; {
+    homepage = "https://github.com/ml-explore/mlx";
+    description = "An array framework for Apple silicon";
+    changelog = "https://github.com/ml-explore/mlx/releases/tag/v${version}";
+    license = licenses.mit;
+    platforms = [ "aarch64-darwin" ];
+    maintainers = with maintainers; [ viraptor ];
+  };
+}

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

@@ -0,0 +1,13 @@
+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)

pkgs/top-level/python-packages.nix

@@ -7173,6 +7173,8 @@ self: super: with self; {
 
   mlrose = callPackage ../development/python-modules/mlrose { };
 
+  mlx = callPackage ../development/python-modules/mlx { };
+
   mlxtend = callPackage ../development/python-modules/mlxtend { };
 
   mlt = toPythonModule (pkgs.mlt.override {