pyrox.dev
1{
2 stdenv,
3 lib,
4 fetchFromGitHub,
5 buildPythonPackage,
6 python3Packages,
7 pybind11,
8 cmake,
9 xcbuild,
10 zsh,
11 darwin,
12 blas,
13 lapack,
14}:
15
16let
17 # static dependencies included directly during compilation
18 gguf-tools = fetchFromGitHub {
19 owner = "antirez";
20 repo = "gguf-tools";
21 rev = "af7d88d808a7608a33723fba067036202910acb3";
22 hash = "sha256-LqNvnUbmq0iziD9VP5OTJCSIy+y/hp5lKCUV7RtKTvM=";
23 };
24 nlohmann_json = fetchFromGitHub {
25 owner = "nlohmann";
26 repo = "json";
27 rev = "v3.11.3";
28 hash = "sha256-7F0Jon+1oWL7uqet5i1IgHX0fUw/+z0QwEcA3zs5xHg=";
29 };
30in
31buildPythonPackage rec {
32 pname = "mlx";
33 version = "0.6.0";
34
35 src = fetchFromGitHub {
36 owner = "ml-explore";
37 repo = "mlx";
38 rev = "refs/tags/v${version}";
39 hash = "sha256-FihdI+3ACKMJfPT2POjTRdtkXs7x+KiQpdpo3RcczBE=";
40 };
41
42 pyproject = true;
43
44 patches = [
45 # With Darwin SDK 11 we cannot include vecLib/cblas_new.h, this needs to wait for PR #229210
46 # In the meantime, pretend Accelerate is not available and use blas/lapack instead.
47 ./disable-accelerate.patch
48 ];
49
50 postPatch = ''
51 substituteInPlace CMakeLists.txt \
52 --replace "/usr/bin/xcrun" "${xcbuild}/bin/xcrun" \
53 '';
54
55 dontUseCmakeConfigure = true;
56
57 env = {
58 PYPI_RELEASE = version;
59 # we can't use Metal compilation with Darwin SDK 11
60 CMAKE_ARGS = toString [
61 (lib.cmakeBool "MLX_BUILD_METAL" false)
62 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
63 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json}")
64 ];
65 };
66
67 nativeBuildInputs = [
68 cmake
69 pybind11
70 xcbuild
71 zsh
72 gguf-tools
73 nlohmann_json
74 ] ++ (with python3Packages; [ setuptools ]);
75
76 buildInputs = [
77 blas
78 lapack
79 ];
80
81 meta = with lib; {
82 homepage = "https://github.com/ml-explore/mlx";
83 description = "An array framework for Apple silicon";
84 changelog = "https://github.com/ml-explore/mlx/releases/tag/v${version}";
85 license = licenses.mit;
86 platforms = [ "aarch64-darwin" ];
87 maintainers = with maintainers; [ viraptor ];
88 };
89}