From 8dde3a1788f1d4a280dfe097178d4fb09c42d7df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stanis=C5=82aw=20Pitucha?= <git@viraptor.info>
Date: Sun, 4 Feb 2024 21:28:05 +1100
Subject: [PATCH] python3Packages.mlx: init at 0.1.0

---
 .../python-modules/mlx/default.nix            | 78 +++++++++++++++++++
 .../mlx/disable-accelerate.patch              | 13 ++++
 pkgs/top-level/python-packages.nix            |  2 +
 3 files changed, 93 insertions(+)
 create mode 100644 pkgs/development/python-modules/mlx/default.nix
 create mode 100644 pkgs/development/python-modules/mlx/disable-accelerate.patch

diff --git a/pkgs/development/python-modules/mlx/default.nix b/pkgs/development/python-modules/mlx/default.nix
new file mode 100644
index 000000000000..036bbfac0418
--- /dev/null
+++ b/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 ];
+  };
+}
diff --git a/pkgs/development/python-modules/mlx/disable-accelerate.patch b/pkgs/development/python-modules/mlx/disable-accelerate.patch
new file mode 100644
index 000000000000..693e7f41104d
--- /dev/null
+++ b/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)
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index f06c38f06d89..d435ab07388b 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -7164,6 +7164,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 {