tjh.dev
1{
2 buildPythonPackage,
3 fetchFromGitHub,
4 lib,
5 replaceVars,
6 stdenv,
7
8 # build-system
9 setuptools,
10
11 # nativeBuildInputs
12 cmake,
13
14 # buildInputs
15 apple-sdk_14,
16 fmt_10,
17 nanobind,
18 nlohmann_json,
19 pybind11,
20
21 # tests
22 numpy,
23 pytestCheckHook,
24 python,
25 runCommand,
26}:
27
28let
29 # static dependencies included directly during compilation
30 gguf-tools = fetchFromGitHub {
31 owner = "antirez";
32 repo = "gguf-tools";
33 rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
34 hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
35 };
36
37 mlx = buildPythonPackage rec {
38 pname = "mlx";
39 version = "0.26.3";
40 pyproject = true;
41
42 src = fetchFromGitHub {
43 owner = "ml-explore";
44 repo = "mlx";
45 tag = "v${version}";
46 hash = "sha256-hbqV/2KYGJ1gyExZd5bgaxTdhl5+Gext+U/+1KAztMU=";
47 };
48
49 patches = [
50 (replaceVars ./darwin-build-fixes.patch {
51 sdkVersion = apple-sdk_14.version;
52 })
53 ];
54
55 postPatch = ''
56 substituteInPlace pyproject.toml \
57 --replace-fail "nanobind==2.4.0" "nanobind>=2.4.0"
58
59 substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
60 --replace-fail "g++" "$CXX"
61 '';
62
63 dontUseCmakeConfigure = true;
64
65 enableParallelBuilding = true;
66
67 # Allows multiple cores to be used in Python builds.
68 postUnpack = ''
69 export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
70 '';
71
72 # updates the wrong fetcher rev attribute
73 passthru.skipBulkUpdate = true;
74
75 env = {
76 PYPI_RELEASE = version;
77 CMAKE_ARGS = toString [
78 # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
79 # To build mlx with Metal support in Nix, you'd need to use one of the sandbox escape
80 # hatches which let you interact with a native install of Xcode, such as `composeXcodeWrapper`
81 # or by changing the upstream (e.g., https://github.com/zed-industries/zed/discussions/7016).
82 (lib.cmakeBool "MLX_BUILD_METAL" false)
83 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
84 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
85 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_FMT" "${fmt_10.src}")
86 ];
87 };
88
89 build-system = [
90 setuptools
91 ];
92
93 nativeBuildInputs = [
94 cmake
95 ];
96
97 buildInputs = [
98 apple-sdk_14
99 fmt_10
100 gguf-tools
101 nanobind
102 nlohmann_json
103 pybind11
104 ];
105
106 pythonImportsCheck = [ "mlx" ];
107
108 # Run the mlx Python test suite.
109 nativeCheckInputs = [
110 numpy
111 pytestCheckHook
112 ];
113
114 enabledTestPaths = [
115 "python/tests/"
116 ];
117
118 # Additional testing by executing the example Python scripts supplied with mlx
119 # using the version of the library we've built.
120 passthru.tests = {
121 mlxTest =
122 runCommand "run-mlx-examples"
123 {
124 buildInputs = [ mlx ];
125 nativeBuildInputs = [ python ];
126 }
127 ''
128 cp ${src}/examples/python/logistic_regression.py .
129 ${python.interpreter} logistic_regression.py
130 rm logistic_regression.py
131
132 cp ${src}/examples/python/linear_regression.py .
133 ${python.interpreter} linear_regression.py
134 rm linear_regression.py
135
136 touch $out
137 '';
138 };
139
140 meta = {
141 homepage = "https://github.com/ml-explore/mlx";
142 description = "Array framework for Apple silicon";
143 changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
144 license = lib.licenses.mit;
145 platforms = [ "aarch64-darwin" ];
146 maintainers = with lib.maintainers; [
147 viraptor
148 Gabriella439
149 cameronyule
150 ];
151 };
152 };
153in
154mlx