tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 replaceVars,
7 nanobind,
8
9 # build-system
10 cmake,
11 setuptools,
12 typing-extensions,
13
14 # buildInputs
15 apple-sdk,
16 fmt,
17 nlohmann_json,
18 # linux-only
19 openblas,
20
21 # tests
22 numpy,
23 pytestCheckHook,
24 python,
25 runCommand,
26
27 # passthru
28 mlx,
29}:
30
31let
32 # static dependencies included directly during compilation
33 gguf-tools = fetchFromGitHub {
34 owner = "antirez";
35 repo = "gguf-tools";
36 rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
37 hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
38 };
39
40in
41buildPythonPackage (finalAttrs: {
42 pname = "mlx";
43 version = "0.30.3";
44 pyproject = true;
45
46 src = fetchFromGitHub {
47 owner = "ml-explore";
48 repo = "mlx";
49 tag = "v${finalAttrs.version}";
50 hash = "sha256-Y4RTkGcDCZ9HLyflN0qYhPt/oVOsBhF1mHnKM4n1/ys=";
51 };
52
53 patches = [
54 # Use nix packages instead of fetching their sources
55 ./dont-fetch-nanobind.patch
56 ./dont-fetch-json.patch
57 ]
58 ++ lib.optionals stdenv.hostPlatform.isDarwin [
59 (replaceVars ./darwin-build-fixes.patch {
60 sdkVersion = apple-sdk.version;
61 })
62 ];
63
64 postPatch = ''
65 substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
66 --replace-fail "g++" "${lib.getExe' stdenv.cc "c++"}"
67 '';
68
69 dontUseCmakeConfigure = true;
70
71 enableParallelBuilding = true;
72
73 # Allows multiple cores to be used in Python builds.
74 postUnpack = ''
75 export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
76 '';
77
78 # updates the wrong fetcher rev attribute
79 passthru.skipBulkUpdate = true;
80
81 env = {
82 PYPI_RELEASE = 1;
83 CMAKE_ARGS = toString (
84 [
85 # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
86 # To build mlx with Metal support in Nix, you'd need to use one of the sandbox escape
87 # hatches which let you interact with a native install of Xcode, such as `composeXcodeWrapper`
88 # or by changing the upstream (e.g., https://github.com/zed-industries/zed/discussions/7016).
89 (lib.cmakeBool "MLX_BUILD_METAL" false)
90 (lib.cmakeBool "USE_SYSTEM_FMT" true)
91 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
92 (lib.cmakeFeature "CMAKE_CXX_FLAGS" "-I${lib.getDev nlohmann_json}/include/nlohmann")
93
94 # Cmake cannot find nanobind-config.cmake by itself
95 (lib.cmakeFeature "nanobind_DIR" "${nanobind}/${python.sitePackages}/nanobind/cmake")
96 ]
97 ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [
98 (lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
99 ]
100 );
101 };
102
103 build-system = [
104 cmake
105 setuptools
106 typing-extensions
107 ];
108
109 buildInputs = [
110 fmt
111 nlohmann_json
112 ]
113 ++ lib.optionals stdenv.hostPlatform.isLinux [
114 openblas
115 ];
116
117 pythonImportsCheck = [ "mlx" ];
118
119 # Run the mlx Python test suite.
120 nativeCheckInputs = [
121 numpy
122 pytestCheckHook
123 ];
124
125 enabledTestPaths = [
126 "python/tests/"
127 ];
128
129 disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [
130 # Segmentation fault
131 "test_lapack"
132 "test_multivariate_normal"
133 "test_orthogonal"
134 "test_vmap_inverse"
135 "test_vmap_svd"
136 ];
137
138 disabledTestPaths = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [
139 # Segmentation fault
140 "python/tests/test_linalg.py"
141 ];
142
143 # patchelf is only available on Linux and no patching is needed on darwin.
144 # Otherwise mlx/core.cpython-313-x86_64-linux-gnu.so contains a reference to
145 # /build/source/build/temp.linux-x86_64-cpython-313/mlx.core/libmlx.so in its rpath.
146 postInstall = lib.optionalString stdenv.hostPlatform.isLinux ''
147 patchelf --replace-needed \
148 libmlx.so \
149 $out/${python.sitePackages}/mlx/lib64/libmlx.so \
150 $out/${python.sitePackages}/mlx/core.cpython-*.so
151 '';
152
153 # Additional testing by executing the example Python scripts supplied with mlx
154 # using the version of the library we've built.
155 passthru.tests = {
156 mlxTest =
157 runCommand "run-mlx-examples"
158 {
159 buildInputs = [ mlx ];
160 nativeBuildInputs = [ python ];
161 }
162 ''
163 cp ${finalAttrs.src}/examples/python/logistic_regression.py .
164 ${python.interpreter} logistic_regression.py
165 rm logistic_regression.py
166
167 cp ${finalAttrs.src}/examples/python/linear_regression.py .
168 ${python.interpreter} linear_regression.py
169 rm linear_regression.py
170
171 touch $out
172 '';
173 };
174
175 meta = {
176 homepage = "https://github.com/ml-explore/mlx";
177 description = "Array framework for Apple silicon";
178 changelog = "https://github.com/ml-explore/mlx/releases/tag/${finalAttrs.src.tag}";
179 license = lib.licenses.mit;
180 maintainers = with lib.maintainers; [
181 Gabriella439
182 booxter
183 cameronyule
184 viraptor
185 ];
186 };
187})